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