| 1 | /* eslint-disable @typescript-eslint/strict-boolean-expressions */ |
| 2 | |
| 3 | import { MessageChannel, SHARE_ENV, Worker, isMainThread } from 'worker_threads' |
| 4 | |
| 5 | function empty (): void {} |
| 6 | const _void = {} |
| 7 | |
| 8 | export type Draft<T> = { -readonly [P in keyof T]?: T[P] } |
| 9 | |
| 10 | export type WorkerWithMessageChannel = Worker & Draft<MessageChannel> |
| 11 | |
| 12 | export interface FixedThreadPoolOptions { |
| 13 | /** |
| 14 | * A function that will listen for error event on each worker thread. |
| 15 | */ |
| 16 | errorHandler?: (this: Worker, e: Error) => void |
| 17 | /** |
| 18 | * A function that will listen for online event on each worker thread. |
| 19 | */ |
| 20 | onlineHandler?: (this: Worker) => void |
| 21 | /** |
| 22 | * A function that will listen for exit event on each worker thread. |
| 23 | */ |
| 24 | exitHandler?: (this: Worker, code: number) => void |
| 25 | /** |
| 26 | * This is just to avoid not useful warnings message, is used to set `maxListeners` on event emitters (workers are event emitters). |
| 27 | * |
| 28 | * @default 1000 |
| 29 | */ |
| 30 | maxTasks?: number |
| 31 | } |
| 32 | |
| 33 | /** |
| 34 | * A thread pool with a static number of threads, is possible to execute tasks in sync or async mode as you prefer. |
| 35 | * |
| 36 | * This pool will select the worker thread in a round robin fashion. |
| 37 | * |
| 38 | * @author [Alessandro Pio Ardizio](https://github.com/pioardi) |
| 39 | * @since 0.0.1 |
| 40 | */ |
| 41 | export default class FixedThreadPool<Data = any, Response = any> { |
| 42 | public readonly workers: WorkerWithMessageChannel[] = [] |
| 43 | public nextWorker: number = 0 |
| 44 | |
| 45 | // threadId as key and an integer value |
| 46 | /* eslint-disable @typescript-eslint/indent */ |
| 47 | public readonly tasks: Map<WorkerWithMessageChannel, number> = new Map< |
| 48 | WorkerWithMessageChannel, |
| 49 | number |
| 50 | >() |
| 51 | /* eslint-enable @typescript-eslint/indent */ |
| 52 | |
| 53 | protected _id: number = 0 |
| 54 | |
| 55 | /** |
| 56 | * @param numThreads Num of threads for this worker pool. |
| 57 | * @param filePath A file path with implementation of `ThreadWorker` class, relative path is fine. |
| 58 | * @param opts An object with possible options for example `errorHandler`, `onlineHandler`. Default: `{ maxTasks: 1000 }` |
| 59 | */ |
| 60 | public constructor ( |
| 61 | public readonly numThreads: number, |
| 62 | public readonly filePath: string, |
| 63 | public readonly opts: FixedThreadPoolOptions = { maxTasks: 1000 } |
| 64 | ) { |
| 65 | if (!isMainThread) { |
| 66 | throw new Error('Cannot start a thread pool from a worker thread !!!') |
| 67 | } |
| 68 | // TODO christopher 2021-02-07: Improve this check e.g. with a pattern or blank check |
| 69 | if (!this.filePath) { |
| 70 | throw new Error('Please specify a file with a worker implementation') |
| 71 | } |
| 72 | |
| 73 | for (let i = 1; i <= this.numThreads; i++) { |
| 74 | this._newWorker() |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | public async destroy (): Promise<void> { |
| 79 | for (const worker of this.workers) { |
| 80 | await worker.terminate() |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | /** |
| 85 | * Execute the task specified into the constructor with the data parameter. |
| 86 | * |
| 87 | * @param data The input for the task specified. |
| 88 | * @returns Promise that is resolved when the task is done. |
| 89 | */ |
| 90 | // eslint-disable-next-line @typescript-eslint/promise-function-async |
| 91 | public execute (data: Data): Promise<Response> { |
| 92 | // configure worker to handle message with the specified task |
| 93 | const worker = this._chooseWorker() |
| 94 | const previousWorkerIndex = this.tasks.get(worker) |
| 95 | if (previousWorkerIndex !== undefined) { |
| 96 | this.tasks.set(worker, previousWorkerIndex + 1) |
| 97 | } else { |
| 98 | throw Error('Worker could not be found in tasks map') |
| 99 | } |
| 100 | const id = ++this._id |
| 101 | const res = this._execute(worker, id) |
| 102 | worker.postMessage({ data: data || _void, _id: id }) |
| 103 | return res |
| 104 | } |
| 105 | |
| 106 | // eslint-disable-next-line @typescript-eslint/promise-function-async |
| 107 | protected _execute ( |
| 108 | worker: WorkerWithMessageChannel, |
| 109 | id: number |
| 110 | ): Promise<Response> { |
| 111 | return new Promise((resolve, reject) => { |
| 112 | const listener = (message: { |
| 113 | _id: number |
| 114 | error?: string |
| 115 | data: Response |
| 116 | }): void => { |
| 117 | if (message._id === id) { |
| 118 | worker.port2?.removeListener('message', listener) |
| 119 | const previousWorkerIndex = this.tasks.get(worker) |
| 120 | if (previousWorkerIndex !== undefined) { |
| 121 | this.tasks.set(worker, previousWorkerIndex + 1) |
| 122 | } else { |
| 123 | throw Error('Worker could not be found in tasks map') |
| 124 | } |
| 125 | if (message.error) reject(message.error) |
| 126 | else resolve(message.data) |
| 127 | } |
| 128 | } |
| 129 | worker.port2?.on('message', listener) |
| 130 | }) |
| 131 | } |
| 132 | |
| 133 | protected _chooseWorker (): WorkerWithMessageChannel { |
| 134 | if (this.workers.length - 1 === this.nextWorker) { |
| 135 | this.nextWorker = 0 |
| 136 | return this.workers[this.nextWorker] |
| 137 | } else { |
| 138 | this.nextWorker++ |
| 139 | return this.workers[this.nextWorker] |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | protected _newWorker (): WorkerWithMessageChannel { |
| 144 | const worker: WorkerWithMessageChannel = new Worker(this.filePath, { |
| 145 | env: SHARE_ENV |
| 146 | }) |
| 147 | worker.on('error', this.opts.errorHandler ?? empty) |
| 148 | worker.on('online', this.opts.onlineHandler ?? empty) |
| 149 | // TODO handle properly when a thread exit |
| 150 | worker.on('exit', this.opts.exitHandler ?? empty) |
| 151 | this.workers.push(worker) |
| 152 | const { port1, port2 } = new MessageChannel() |
| 153 | worker.postMessage({ parent: port1 }, [port1]) |
| 154 | worker.port1 = port1 |
| 155 | worker.port2 = port2 |
| 156 | // we will attach a listener for every task, |
| 157 | // when task is completed the listener will be removed but to avoid warnings we are increasing the max listeners size |
| 158 | worker.port2.setMaxListeners(this.opts.maxTasks ?? 1000) |
| 159 | // init tasks map |
| 160 | this.tasks.set(worker, 0) |
| 161 | return worker |
| 162 | } |
| 163 | } |
| 164 | |
| 165 | module.exports = FixedThreadPool |