95 lines
2.4 KiB
JavaScript
95 lines
2.4 KiB
JavaScript
import { Worker, parentPort } from 'worker_threads';
|
|
import { TransformStream } from 'web-streams-polyfill';
|
|
|
|
function uuid() {
|
|
return Array.from({ length: 16 }, () =>
|
|
Math.floor(Math.random() * 256).toString(16),
|
|
).join('');
|
|
}
|
|
|
|
function jobPromise(worker, msg) {
|
|
return new Promise((resolve) => {
|
|
const id = uuid();
|
|
worker.postMessage({ msg, id });
|
|
worker.on('message', function f({ result, id: rid }) {
|
|
if (rid !== id) {
|
|
return;
|
|
}
|
|
worker.off('message', f);
|
|
resolve(result);
|
|
});
|
|
worker.on('error', (error) => console.error('Worker error: ', error));
|
|
});
|
|
}
|
|
|
|
export default class WorkerPool {
|
|
constructor(numWorkers, workerFile) {
|
|
this.numWorkers = numWorkers;
|
|
this.jobQueue = new TransformStream();
|
|
this.workerQueue = new TransformStream();
|
|
|
|
const writer = this.workerQueue.writable.getWriter();
|
|
for (let i = 0; i < numWorkers; i++) {
|
|
writer.write(new Worker(workerFile));
|
|
}
|
|
writer.releaseLock();
|
|
|
|
this.done = this._readLoop();
|
|
}
|
|
|
|
async _readLoop() {
|
|
const reader = this.jobQueue.readable.getReader();
|
|
while (true) {
|
|
const { value, done } = await reader.read();
|
|
if (done) {
|
|
await this._terminateAll();
|
|
return;
|
|
}
|
|
const { msg, resolve } = value;
|
|
const worker = await this._nextWorker();
|
|
jobPromise(worker, msg).then((result) => {
|
|
resolve(result);
|
|
const writer = this.workerQueue.writable.getWriter();
|
|
writer.write(worker);
|
|
writer.releaseLock();
|
|
});
|
|
}
|
|
}
|
|
|
|
async _nextWorker() {
|
|
const reader = this.workerQueue.readable.getReader();
|
|
const { value, done } = await reader.read();
|
|
reader.releaseLock();
|
|
return value;
|
|
}
|
|
|
|
async _terminateAll() {
|
|
for (let n = 0; n < this.numWorkers; n++) {
|
|
const worker = await this._nextWorker();
|
|
worker.terminate();
|
|
}
|
|
this.workerQueue.writable.close();
|
|
}
|
|
|
|
async join() {
|
|
this.jobQueue.writable.getWriter().close();
|
|
await this.done;
|
|
}
|
|
|
|
dispatchJob(msg) {
|
|
return new Promise((resolve) => {
|
|
const writer = this.jobQueue.writable.getWriter();
|
|
writer.write({ msg, resolve });
|
|
writer.releaseLock();
|
|
});
|
|
}
|
|
|
|
static useThisThreadAsWorker(cb) {
|
|
parentPort.on('message', async (data) => {
|
|
const { msg, id } = data;
|
|
const result = await cb(msg);
|
|
parentPort.postMessage({ result, id });
|
|
});
|
|
}
|
|
}
|