forked from microsoft/rushstack
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathWorkerPool.ts
More file actions
264 lines (225 loc) · 6.56 KB
/
Copy pathWorkerPool.ts
File metadata and controls
264 lines (225 loc) · 6.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
// Copyright (c) Microsoft Corporation. All rights reserved. Licensed under the MIT license.
// See LICENSE in the project root for license information.
import { Worker } from 'worker_threads';
/**
* Symbol to read the ID off of a worker
* @internal
*/
export const WORKER_ID_SYMBOL: unique symbol = Symbol('workerId');
/**
* @internal
*/
export interface IWorkerPoolOptions {
/**
* Identifier for this pool, to assign to its workers for tracking
*/
id: string;
/**
* Maximum number of concurrent workers this WorkerPool may spawn.
*/
maxWorkers: number;
/**
* Optional callback invoked when a worker is destroyed.
*/
onWorkerDestroyed?: () => void;
/**
* Optional callback invoked on a newly created worker.
*/
prepareWorker?: (worker: Worker) => void;
/**
* Optional data to pass to workers when they are initialized.
* Will be subjected to the Structured Clone algorithm.
*/
workerData?: unknown;
/**
* Absolute path to the worker script.
*/
workerScriptPath: string;
}
/**
* Manages a pool of workers.
* Workers will be shutdown by sending them the boolean value `false` in a postMessage.
* @internal
*/
export class WorkerPool {
public id: string;
public maxWorkers: number;
private readonly _alive: Worker[];
private _error: Error | undefined;
private _finishing: boolean;
private readonly _idle: Worker[];
private _nextId: number;
private readonly _onComplete: [() => void, (error: Error) => void][];
private readonly _onWorkerDestroyed: (() => void) | undefined;
private readonly _pending: [(worker: Worker) => void, (error: Error) => void][];
private readonly _prepare: ((worker: Worker) => void) | undefined;
private readonly _workerData: unknown;
private readonly _workerScript: string;
public constructor(options: IWorkerPoolOptions) {
const { id, maxWorkers, onWorkerDestroyed, prepareWorker, workerData, workerScriptPath } = options;
this.id = id;
this.maxWorkers = maxWorkers;
this._alive = [];
this._error = undefined;
this._finishing = false;
this._idle = [];
this._nextId = 0;
this._onComplete = [];
this._onWorkerDestroyed = onWorkerDestroyed;
this._pending = [];
this._prepare = prepareWorker;
this._workerData = workerData;
this._workerScript = workerScriptPath;
}
/**
* Gets the count of active workers.
*/
public getActiveCount(): number {
return this._alive.length - this._idle.length;
}
/**
* Gets the count of idle workers.
*/
public getIdleCount(): number {
return this._idle.length;
}
/**
* Gets the count of live workers.
*/
public getLiveCount(): number {
return this._alive.length;
}
/**
* Tells the pool to shut down when all workers are done.
* Returns a promise that will be fulfilled if all workers finish successfully, or reject with the first error.
*/
public async finishAsync(): Promise<void> {
this._finishing = true;
if (this._error) {
throw this._error;
}
if (!this._alive.length) {
// The pool has no live workers, this is a no-op
return;
}
// Clean up all idle workers
for (const worker of this._idle.splice(0)) {
worker.postMessage(false);
}
// There are still active workers, wait for them to clean up.
await new Promise<void>((resolve, reject) => this._onComplete.push([resolve, reject]));
}
/**
* Resets the pool and allows more work
*/
public reset(): void {
this._finishing = false;
this._error = undefined;
}
/**
* Returns a worker to the pool. If the pool is finishing, deallocates the worker.
* @param worker - The worker to free
*/
public checkinWorker(worker: Worker): void {
if (this._error) {
// Shut down the worker (failure)
worker.postMessage(false);
return;
}
const next: [(worker: Worker) => void, unknown] | undefined = this._pending.shift();
if (next) {
// Perform the next unit of work;
next[0](worker);
} else if (this._finishing) {
// Shut down the worker (success)
worker.postMessage(false);
} else {
// No pending work, idle the workers
this._idle.push(worker);
}
}
/**
* Checks out a currently available worker or waits for the next free worker.
* @param allowCreate - If creating new workers is allowed (subject to maxSize)
*/
public async checkoutWorkerAsync(allowCreate: boolean): Promise<Worker> {
if (this._error) {
throw this._error;
}
let worker: Worker | undefined = this._idle.shift();
if (!worker && allowCreate) {
worker = this._createWorker();
}
if (worker) {
return worker;
}
return await new Promise((resolve: (worker: Worker) => void, reject: (error: Error) => void) => {
this._pending.push([resolve, reject]);
});
}
/**
* Creates a new worker if allowed by maxSize.
*/
private _createWorker(): Worker | undefined {
if (this._alive.length >= this.maxWorkers) {
return;
}
const worker: Worker & {
[WORKER_ID_SYMBOL]?: string;
} = new Worker(this._workerScript, {
eval: false,
workerData: this._workerData
});
const id: string = `${this.id}#${++this._nextId}`;
worker[WORKER_ID_SYMBOL] = id;
this._alive.push(worker);
worker.on('error', (err) => {
this._onError(err);
this._destroyWorker(worker);
});
worker.on('exit', (exitCode) => {
if (exitCode !== 0) {
this._onError(new Error(`Worker ${id} exited with code ${exitCode}`));
}
this._destroyWorker(worker);
});
if (this._prepare) {
this._prepare(worker);
}
return worker;
}
/**
* Cleans up a worker
*/
private _destroyWorker(worker: Worker): void {
const aliveIndex: number = this._alive.indexOf(worker);
if (aliveIndex >= 0) {
this._alive.splice(aliveIndex, 1);
}
const freeIndex: number = this._idle.indexOf(worker);
if (freeIndex >= 0) {
this._idle.splice(freeIndex, 1);
}
worker.unref();
if (this._onWorkerDestroyed) {
this._onWorkerDestroyed();
}
if (!this._alive.length && !this._error) {
for (const [resolve] of this._onComplete.splice(0)) {
resolve();
}
}
}
/**
* Notifies all pending callbacks that an error has occurred and switches this pool into error state.
*/
private _onError(error: Error): void {
this._error = error;
for (const [, reject] of this._pending.splice(0)) {
reject(this._error);
}
for (const [, reject] of this._onComplete.splice(0)) {
reject(this._error);
}
}
}