Spaces:
Sleeping
Sleeping
File size: 3,053 Bytes
391a73c | 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 | import { AsyncResource } from 'async_hooks';
import { Worker } from 'worker_threads';
import { cpus } from 'os';
import { EventEmitter } from 'events';
import serializeJavascript from 'serialize-javascript';
import { freeWorker, taskInfo, workerPoolWorkerFlag } from './constants';
import type {
WorkerCallback,
WorkerContext,
WorkerOutput,
WorkerPoolOptions,
WorkerPoolTask,
WorkerWithTaskInfo
} from './type';
class WorkerPoolTaskInfo extends AsyncResource {
constructor(private callback: WorkerCallback) {
super('WorkerPoolTaskInfo');
}
done(err: Error | null, result: any) {
this.runInAsyncScope(this.callback, null, err, result);
this.emitDestroy();
}
}
export class WorkerPool extends EventEmitter {
protected maxInstances: number;
protected filePath: string;
protected tasks: WorkerPoolTask[] = [];
protected workers: WorkerWithTaskInfo[] = [];
protected freeWorkers: WorkerWithTaskInfo[] = [];
constructor(options: WorkerPoolOptions) {
super();
this.maxInstances = options.maxWorkers || cpus().length;
this.filePath = options.filePath;
this.on(freeWorker, () => {
if (this.tasks.length > 0) {
const { context, cb } = this.tasks.shift()!;
this.runTask(context, cb);
}
});
}
get numWorkers(): number {
return this.workers.length;
}
addAsync(context: WorkerContext): Promise<WorkerOutput> {
return new Promise((resolve, reject) => {
this.runTask(context, (err, output) => {
if (err) {
reject(err);
return;
}
if (!output) {
reject(new Error('The output is empty'));
return;
}
resolve(output);
});
});
}
close() {
for (let i = 0; i < this.workers.length; i++) {
const worker = this.workers[i];
worker.terminate();
}
}
private addNewWorker() {
const worker: WorkerWithTaskInfo = new Worker(this.filePath, {
workerData: workerPoolWorkerFlag
});
worker.on('message', (result) => {
worker[taskInfo]?.done(null, result);
worker[taskInfo] = null;
this.freeWorkers.push(worker);
this.emit(freeWorker);
});
worker.on('error', (err) => {
if (worker[taskInfo]) {
worker[taskInfo].done(err, null);
} else {
this.emit('error', err);
}
this.workers.splice(this.workers.indexOf(worker), 1);
this.addNewWorker();
});
this.workers.push(worker);
this.freeWorkers.push(worker);
this.emit(freeWorker);
}
private runTask(context: WorkerContext, cb: WorkerCallback) {
if (this.freeWorkers.length === 0) {
this.tasks.push({ context, cb });
if (this.numWorkers < this.maxInstances) {
this.addNewWorker();
}
return;
}
const worker = this.freeWorkers.pop();
if (worker) {
worker[taskInfo] = new WorkerPoolTaskInfo(cb);
worker.postMessage({
code: context.code,
options: serializeJavascript(context.options)
});
}
}
}
|