Merge pull request #747 from poolifier/multiple-functions
[poolifier.git] / src / pools / abstract-pool.ts
1 import crypto from 'node:crypto'
2 import type { MessageValue, PromiseResponseWrapper } from '../utility-types'
3 import {
4 DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS,
5 EMPTY_FUNCTION,
6 median
7 } from '../utils'
8 import { KillBehaviors, isKillBehavior } from '../worker/worker-options'
9 import { CircularArray } from '../circular-array'
10 import { Queue } from '../queue'
11 import {
12 type IPool,
13 PoolEmitter,
14 PoolEvents,
15 type PoolOptions,
16 PoolType,
17 type TasksQueueOptions
18 } from './pool'
19 import type { IWorker, Task, TasksUsage, WorkerNode } from './worker'
20 import {
21 WorkerChoiceStrategies,
22 type WorkerChoiceStrategy,
23 type WorkerChoiceStrategyOptions
24 } from './selection-strategies/selection-strategies-types'
25 import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
26
27 /**
28 * Base class that implements some shared logic for all poolifier pools.
29 *
30 * @typeParam Worker - Type of worker which manages this pool.
31 * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
32 * @typeParam Response - Type of execution response. This can only be serializable data.
33 */
34 export abstract class AbstractPool<
35 Worker extends IWorker,
36 Data = unknown,
37 Response = unknown
38 > implements IPool<Worker, Data, Response> {
39 /** @inheritDoc */
40 public readonly workerNodes: Array<WorkerNode<Worker, Data>> = []
41
42 /** @inheritDoc */
43 public readonly emitter?: PoolEmitter
44
45 /**
46 * The execution response promise map.
47 *
48 * - `key`: The message id of each submitted task.
49 * - `value`: An object that contains the worker, the execution response promise resolve and reject callbacks.
50 *
51 * When we receive a message from the worker, we get a map entry with the promise resolve/reject bound to the message id.
52 */
53 protected promiseResponseMap: Map<
54 string,
55 PromiseResponseWrapper<Worker, Response>
56 > = new Map<string, PromiseResponseWrapper<Worker, Response>>()
57
58 /**
59 * Worker choice strategy context referencing a worker choice algorithm implementation.
60 *
61 * Default to a round robin algorithm.
62 */
63 protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
64 Worker,
65 Data,
66 Response
67 >
68
69 /**
70 * Constructs a new poolifier pool.
71 *
72 * @param numberOfWorkers - Number of workers that this pool should manage.
73 * @param filePath - Path to the worker file.
74 * @param opts - Options for the pool.
75 */
76 public constructor (
77 public readonly numberOfWorkers: number,
78 public readonly filePath: string,
79 public readonly opts: PoolOptions<Worker>
80 ) {
81 if (!this.isMain()) {
82 throw new Error('Cannot start a pool from a worker!')
83 }
84 this.checkNumberOfWorkers(this.numberOfWorkers)
85 this.checkFilePath(this.filePath)
86 this.checkPoolOptions(this.opts)
87
88 this.chooseWorkerNode = this.chooseWorkerNode.bind(this)
89 this.executeTask = this.executeTask.bind(this)
90 this.enqueueTask = this.enqueueTask.bind(this)
91 this.checkAndEmitEvents = this.checkAndEmitEvents.bind(this)
92
93 this.setupHook()
94
95 for (let i = 1; i <= this.numberOfWorkers; i++) {
96 this.createAndSetupWorker()
97 }
98
99 if (this.opts.enableEvents === true) {
100 this.emitter = new PoolEmitter()
101 }
102 this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext<
103 Worker,
104 Data,
105 Response
106 >(
107 this,
108 this.opts.workerChoiceStrategy,
109 this.opts.workerChoiceStrategyOptions
110 )
111 }
112
113 private checkFilePath (filePath: string): void {
114 if (
115 filePath == null ||
116 (typeof filePath === 'string' && filePath.trim().length === 0)
117 ) {
118 throw new Error('Please specify a file with a worker implementation')
119 }
120 }
121
122 private checkNumberOfWorkers (numberOfWorkers: number): void {
123 if (numberOfWorkers == null) {
124 throw new Error(
125 'Cannot instantiate a pool without specifying the number of workers'
126 )
127 } else if (!Number.isSafeInteger(numberOfWorkers)) {
128 throw new TypeError(
129 'Cannot instantiate a pool with a non integer number of workers'
130 )
131 } else if (numberOfWorkers < 0) {
132 throw new RangeError(
133 'Cannot instantiate a pool with a negative number of workers'
134 )
135 } else if (this.type === PoolType.FIXED && numberOfWorkers === 0) {
136 throw new Error('Cannot instantiate a fixed pool with no worker')
137 }
138 }
139
140 private checkPoolOptions (opts: PoolOptions<Worker>): void {
141 this.opts.workerChoiceStrategy =
142 opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
143 this.checkValidWorkerChoiceStrategy(this.opts.workerChoiceStrategy)
144 this.opts.workerChoiceStrategyOptions =
145 opts.workerChoiceStrategyOptions ?? DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
146 this.opts.enableEvents = opts.enableEvents ?? true
147 this.opts.enableTasksQueue = opts.enableTasksQueue ?? false
148 if (this.opts.enableTasksQueue) {
149 this.checkValidTasksQueueOptions(
150 opts.tasksQueueOptions as TasksQueueOptions
151 )
152 this.opts.tasksQueueOptions = this.buildTasksQueueOptions(
153 opts.tasksQueueOptions as TasksQueueOptions
154 )
155 }
156 }
157
158 private checkValidWorkerChoiceStrategy (
159 workerChoiceStrategy: WorkerChoiceStrategy
160 ): void {
161 if (!Object.values(WorkerChoiceStrategies).includes(workerChoiceStrategy)) {
162 throw new Error(
163 `Invalid worker choice strategy '${workerChoiceStrategy}'`
164 )
165 }
166 }
167
168 private checkValidTasksQueueOptions (
169 tasksQueueOptions: TasksQueueOptions
170 ): void {
171 if ((tasksQueueOptions?.concurrency as number) <= 0) {
172 throw new Error(
173 `Invalid worker tasks concurrency '${
174 tasksQueueOptions.concurrency as number
175 }'`
176 )
177 }
178 }
179
180 /** @inheritDoc */
181 public abstract get type (): PoolType
182
183 /**
184 * Number of tasks running in the pool.
185 */
186 private get numberOfRunningTasks (): number {
187 return this.workerNodes.reduce(
188 (accumulator, workerNode) => accumulator + workerNode.tasksUsage.running,
189 0
190 )
191 }
192
193 /**
194 * Number of tasks queued in the pool.
195 */
196 private get numberOfQueuedTasks (): number {
197 if (this.opts.enableTasksQueue === false) {
198 return 0
199 }
200 return this.workerNodes.reduce(
201 (accumulator, workerNode) => accumulator + workerNode.tasksQueue.size,
202 0
203 )
204 }
205
206 /**
207 * Gets the given worker its worker node key.
208 *
209 * @param worker - The worker.
210 * @returns The worker node key if the worker is found in the pool worker nodes, `-1` otherwise.
211 */
212 private getWorkerNodeKey (worker: Worker): number {
213 return this.workerNodes.findIndex(
214 workerNode => workerNode.worker === worker
215 )
216 }
217
218 /** @inheritDoc */
219 public setWorkerChoiceStrategy (
220 workerChoiceStrategy: WorkerChoiceStrategy,
221 workerChoiceStrategyOptions?: WorkerChoiceStrategyOptions
222 ): void {
223 this.checkValidWorkerChoiceStrategy(workerChoiceStrategy)
224 this.opts.workerChoiceStrategy = workerChoiceStrategy
225 for (const workerNode of this.workerNodes) {
226 this.setWorkerNodeTasksUsage(workerNode, {
227 run: 0,
228 running: 0,
229 runTime: 0,
230 runTimeHistory: new CircularArray(),
231 avgRunTime: 0,
232 medRunTime: 0,
233 error: 0
234 })
235 }
236 this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
237 this.opts.workerChoiceStrategy
238 )
239 if (workerChoiceStrategyOptions != null) {
240 this.setWorkerChoiceStrategyOptions(workerChoiceStrategyOptions)
241 }
242 }
243
244 /** @inheritDoc */
245 public setWorkerChoiceStrategyOptions (
246 workerChoiceStrategyOptions: WorkerChoiceStrategyOptions
247 ): void {
248 this.opts.workerChoiceStrategyOptions = workerChoiceStrategyOptions
249 this.workerChoiceStrategyContext.setOptions(
250 this.opts.workerChoiceStrategyOptions
251 )
252 }
253
254 /** @inheritDoc */
255 public enableTasksQueue (
256 enable: boolean,
257 tasksQueueOptions?: TasksQueueOptions
258 ): void {
259 if (this.opts.enableTasksQueue === true && !enable) {
260 this.flushTasksQueues()
261 }
262 this.opts.enableTasksQueue = enable
263 this.setTasksQueueOptions(tasksQueueOptions as TasksQueueOptions)
264 }
265
266 /** @inheritDoc */
267 public setTasksQueueOptions (tasksQueueOptions: TasksQueueOptions): void {
268 if (this.opts.enableTasksQueue === true) {
269 this.checkValidTasksQueueOptions(tasksQueueOptions)
270 this.opts.tasksQueueOptions =
271 this.buildTasksQueueOptions(tasksQueueOptions)
272 } else {
273 delete this.opts.tasksQueueOptions
274 }
275 }
276
277 private buildTasksQueueOptions (
278 tasksQueueOptions: TasksQueueOptions
279 ): TasksQueueOptions {
280 return {
281 concurrency: tasksQueueOptions?.concurrency ?? 1
282 }
283 }
284
285 /**
286 * Whether the pool is full or not.
287 *
288 * The pool filling boolean status.
289 */
290 protected abstract get full (): boolean
291
292 /**
293 * Whether the pool is busy or not.
294 *
295 * The pool busyness boolean status.
296 */
297 protected abstract get busy (): boolean
298
299 protected internalBusy (): boolean {
300 return (
301 this.workerNodes.findIndex(workerNode => {
302 return workerNode.tasksUsage?.running === 0
303 }) === -1
304 )
305 }
306
307 /** @inheritDoc */
308 public async execute (data?: Data, name?: string): Promise<Response> {
309 const [workerNodeKey, workerNode] = this.chooseWorkerNode()
310 const submittedTask: Task<Data> = {
311 name,
312 // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
313 data: data ?? ({} as Data),
314 id: crypto.randomUUID()
315 }
316 const res = new Promise<Response>((resolve, reject) => {
317 this.promiseResponseMap.set(submittedTask.id as string, {
318 resolve,
319 reject,
320 worker: workerNode.worker
321 })
322 })
323 if (
324 this.opts.enableTasksQueue === true &&
325 (this.busy ||
326 this.workerNodes[workerNodeKey].tasksUsage.running >=
327 ((this.opts.tasksQueueOptions as TasksQueueOptions)
328 .concurrency as number))
329 ) {
330 this.enqueueTask(workerNodeKey, submittedTask)
331 } else {
332 this.executeTask(workerNodeKey, submittedTask)
333 }
334 this.checkAndEmitEvents()
335 // eslint-disable-next-line @typescript-eslint/return-await
336 return res
337 }
338
339 /** @inheritDoc */
340 public async destroy (): Promise<void> {
341 await Promise.all(
342 this.workerNodes.map(async (workerNode, workerNodeKey) => {
343 this.flushTasksQueue(workerNodeKey)
344 await this.destroyWorker(workerNode.worker)
345 })
346 )
347 }
348
349 /**
350 * Shutdowns the given worker.
351 *
352 * @param worker - A worker within `workerNodes`.
353 */
354 protected abstract destroyWorker (worker: Worker): void | Promise<void>
355
356 /**
357 * Setup hook to execute code before worker node are created in the abstract constructor.
358 * Can be overridden
359 *
360 * @virtual
361 */
362 protected setupHook (): void {
363 // Intentionally empty
364 }
365
366 /**
367 * Should return whether the worker is the main worker or not.
368 */
369 protected abstract isMain (): boolean
370
371 /**
372 * Hook executed before the worker task execution.
373 * Can be overridden.
374 *
375 * @param workerNodeKey - The worker node key.
376 */
377 protected beforeTaskExecutionHook (workerNodeKey: number): void {
378 ++this.workerNodes[workerNodeKey].tasksUsage.running
379 }
380
381 /**
382 * Hook executed after the worker task execution.
383 * Can be overridden.
384 *
385 * @param worker - The worker.
386 * @param message - The received message.
387 */
388 protected afterTaskExecutionHook (
389 worker: Worker,
390 message: MessageValue<Response>
391 ): void {
392 const workerTasksUsage = this.getWorkerTasksUsage(worker)
393 --workerTasksUsage.running
394 ++workerTasksUsage.run
395 if (message.error != null) {
396 ++workerTasksUsage.error
397 }
398 if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
399 workerTasksUsage.runTime += message.runTime ?? 0
400 if (
401 this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
402 workerTasksUsage.run !== 0
403 ) {
404 workerTasksUsage.avgRunTime =
405 workerTasksUsage.runTime / workerTasksUsage.run
406 }
407 if (this.workerChoiceStrategyContext.getRequiredStatistics().medRunTime) {
408 workerTasksUsage.runTimeHistory.push(message.runTime ?? 0)
409 workerTasksUsage.medRunTime = median(workerTasksUsage.runTimeHistory)
410 }
411 }
412 }
413
414 /**
415 * Chooses a worker node for the next task.
416 *
417 * The default uses a round robin algorithm to distribute the load.
418 *
419 * @returns [worker node key, worker node].
420 */
421 protected chooseWorkerNode (): [number, WorkerNode<Worker, Data>] {
422 let workerNodeKey: number
423 if (this.type === PoolType.DYNAMIC && !this.full && this.internalBusy()) {
424 const workerCreated = this.createAndSetupWorker()
425 this.registerWorkerMessageListener(workerCreated, message => {
426 if (
427 isKillBehavior(KillBehaviors.HARD, message.kill) ||
428 (message.kill != null &&
429 this.getWorkerTasksUsage(workerCreated)?.running === 0)
430 ) {
431 // Kill message received from the worker: no new tasks are submitted to that worker for a while ( > maxInactiveTime)
432 this.flushTasksQueueByWorker(workerCreated)
433 void (this.destroyWorker(workerCreated) as Promise<void>)
434 }
435 })
436 workerNodeKey = this.getWorkerNodeKey(workerCreated)
437 } else {
438 workerNodeKey = this.workerChoiceStrategyContext.execute()
439 }
440 return [workerNodeKey, this.workerNodes[workerNodeKey]]
441 }
442
443 /**
444 * Sends a message to the given worker.
445 *
446 * @param worker - The worker which should receive the message.
447 * @param message - The message.
448 */
449 protected abstract sendToWorker (
450 worker: Worker,
451 message: MessageValue<Data>
452 ): void
453
454 /**
455 * Registers a listener callback on the given worker.
456 *
457 * @param worker - The worker which should register a listener.
458 * @param listener - The message listener callback.
459 */
460 protected abstract registerWorkerMessageListener<
461 Message extends Data | Response
462 >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
463
464 /**
465 * Returns a newly created worker.
466 */
467 protected abstract createWorker (): Worker
468
469 /**
470 * Function that can be hooked up when a worker has been newly created and moved to the pool worker nodes.
471 *
472 * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
473 *
474 * @param worker - The newly created worker.
475 */
476 protected abstract afterWorkerSetup (worker: Worker): void
477
478 /**
479 * Creates a new worker and sets it up completely in the pool worker nodes.
480 *
481 * @returns New, completely set up worker.
482 */
483 protected createAndSetupWorker (): Worker {
484 const worker = this.createWorker()
485
486 worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
487 worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
488 worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
489 worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
490 worker.once('exit', () => {
491 this.removeWorkerNode(worker)
492 })
493
494 this.pushWorkerNode(worker)
495
496 this.afterWorkerSetup(worker)
497
498 return worker
499 }
500
501 /**
502 * This function is the listener registered for each worker message.
503 *
504 * @returns The listener function to execute when a message is received from a worker.
505 */
506 protected workerListener (): (message: MessageValue<Response>) => void {
507 return message => {
508 if (message.id != null) {
509 // Task execution response received
510 const promiseResponse = this.promiseResponseMap.get(message.id)
511 if (promiseResponse != null) {
512 if (message.error != null) {
513 promiseResponse.reject(message.error)
514 } else {
515 promiseResponse.resolve(message.data as Response)
516 }
517 this.afterTaskExecutionHook(promiseResponse.worker, message)
518 this.promiseResponseMap.delete(message.id)
519 const workerNodeKey = this.getWorkerNodeKey(promiseResponse.worker)
520 if (
521 this.opts.enableTasksQueue === true &&
522 this.tasksQueueSize(workerNodeKey) > 0
523 ) {
524 this.executeTask(
525 workerNodeKey,
526 this.dequeueTask(workerNodeKey) as Task<Data>
527 )
528 }
529 }
530 }
531 }
532 }
533
534 private checkAndEmitEvents (): void {
535 if (this.opts.enableEvents === true) {
536 if (this.busy) {
537 this.emitter?.emit(PoolEvents.busy)
538 }
539 if (this.type === PoolType.DYNAMIC && this.full) {
540 this.emitter?.emit(PoolEvents.full)
541 }
542 }
543 }
544
545 /**
546 * Sets the given worker node its tasks usage in the pool.
547 *
548 * @param workerNode - The worker node.
549 * @param tasksUsage - The worker node tasks usage.
550 */
551 private setWorkerNodeTasksUsage (
552 workerNode: WorkerNode<Worker, Data>,
553 tasksUsage: TasksUsage
554 ): void {
555 workerNode.tasksUsage = tasksUsage
556 }
557
558 /**
559 * Gets the given worker its tasks usage in the pool.
560 *
561 * @param worker - The worker.
562 * @throws Error if the worker is not found in the pool worker nodes.
563 * @returns The worker tasks usage.
564 */
565 private getWorkerTasksUsage (worker: Worker): TasksUsage {
566 const workerNodeKey = this.getWorkerNodeKey(worker)
567 if (workerNodeKey !== -1) {
568 return this.workerNodes[workerNodeKey].tasksUsage
569 }
570 throw new Error('Worker could not be found in the pool worker nodes')
571 }
572
573 /**
574 * Pushes the given worker in the pool worker nodes.
575 *
576 * @param worker - The worker.
577 * @returns The worker nodes length.
578 */
579 private pushWorkerNode (worker: Worker): number {
580 return this.workerNodes.push({
581 worker,
582 tasksUsage: {
583 run: 0,
584 running: 0,
585 runTime: 0,
586 runTimeHistory: new CircularArray(),
587 avgRunTime: 0,
588 medRunTime: 0,
589 error: 0
590 },
591 tasksQueue: new Queue<Task<Data>>()
592 })
593 }
594
595 /**
596 * Sets the given worker in the pool worker nodes.
597 *
598 * @param workerNodeKey - The worker node key.
599 * @param worker - The worker.
600 * @param tasksUsage - The worker tasks usage.
601 * @param tasksQueue - The worker task queue.
602 */
603 private setWorkerNode (
604 workerNodeKey: number,
605 worker: Worker,
606 tasksUsage: TasksUsage,
607 tasksQueue: Queue<Task<Data>>
608 ): void {
609 this.workerNodes[workerNodeKey] = {
610 worker,
611 tasksUsage,
612 tasksQueue
613 }
614 }
615
616 /**
617 * Removes the given worker from the pool worker nodes.
618 *
619 * @param worker - The worker.
620 */
621 private removeWorkerNode (worker: Worker): void {
622 const workerNodeKey = this.getWorkerNodeKey(worker)
623 this.workerNodes.splice(workerNodeKey, 1)
624 this.workerChoiceStrategyContext.remove(workerNodeKey)
625 }
626
627 private executeTask (workerNodeKey: number, task: Task<Data>): void {
628 this.beforeTaskExecutionHook(workerNodeKey)
629 this.sendToWorker(this.workerNodes[workerNodeKey].worker, task)
630 }
631
632 private enqueueTask (workerNodeKey: number, task: Task<Data>): number {
633 return this.workerNodes[workerNodeKey].tasksQueue.enqueue(task)
634 }
635
636 private dequeueTask (workerNodeKey: number): Task<Data> | undefined {
637 return this.workerNodes[workerNodeKey].tasksQueue.dequeue()
638 }
639
640 private tasksQueueSize (workerNodeKey: number): number {
641 return this.workerNodes[workerNodeKey].tasksQueue.size
642 }
643
644 private flushTasksQueue (workerNodeKey: number): void {
645 if (this.tasksQueueSize(workerNodeKey) > 0) {
646 for (let i = 0; i < this.tasksQueueSize(workerNodeKey); i++) {
647 this.executeTask(
648 workerNodeKey,
649 this.dequeueTask(workerNodeKey) as Task<Data>
650 )
651 }
652 }
653 }
654
655 private flushTasksQueueByWorker (worker: Worker): void {
656 const workerNodeKey = this.getWorkerNodeKey(worker)
657 this.flushTasksQueue(workerNodeKey)
658 }
659
660 private flushTasksQueues (): void {
661 for (const [workerNodeKey] of this.workerNodes.entries()) {
662 this.flushTasksQueue(workerNodeKey)
663 }
664 }
665 }