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