with-worker-threads/dist/index.js

Spawn worker threads that are closed when the function returns

import { isMainThread, Worker, parentPort } from "node:worker_threads";
import { availableParallelism } from "node:os";
const makeWorkerPool = (path, options) => {
    const closed = new AbortController();
    const workers = Array(options?.concurrency ?? availableParallelism()).fill(null).map(() => {
        const worker = new Worker(path, options?.workerOptions);
        return { worker, tasks: [] };
    });
    closed.signal.addEventListener("abort", () => {
        workers.forEach(({ worker }) => {
            worker.postMessage({ close: true });
        });
    });
    const withResolvers = () => {
        let resolve, reject;
        const promise = new Promise((res, rej) => {
            resolve = res;
            reject = rej;
        });
        return { promise, resolve: resolve, reject: reject };
    };
    const pending = [];
    const runTask = async (workerObj, fn, resolve, reject) => {
        const task = Symbol();
        workerObj.tasks.push(task);
        await Promise.resolve(fn(workerObj.worker)).then(resolve, reject).catch(() => { });
        workerObj.tasks.splice(workerObj.tasks.indexOf(task), 1);
        if (pending.length > 0) {
            const nextTask = pending.pop();
            await runTask(workerObj, nextTask.task, nextTask.resolve, nextTask.reject);
        }
    };
    const abortQueue = new Map();
    closed.signal.addEventListener("abort", () => {
        abortQueue.forEach((fn) => fn());
    });
    const postTask = async (fn) => {
        if (closed.signal.aborted) {
            return Promise.reject(closed.signal.reason);
        }
        else {
            const { promise, resolve, reject } = withResolvers();
            const abortSymbol = Symbol();
            abortQueue.set(abortSymbol, () => reject(closed.signal.reason));
            promise.catch(() => { }).then(() => abortQueue.delete(abortSymbol));
            const availableWorker = workers.filter(({ tasks }) => tasks.length < (options?.maxUtilization ?? 1)).sort((a, b) => a.tasks.length - b.tasks.length)[0];
            if (availableWorker) {
                runTask(availableWorker, fn, resolve, reject);
            }
            else {
                pending.push({ task: fn, resolve, reject });
            }
            return promise;
        }
    };
    return {
        task: (operation) => async (args, transferList) => {
            return postTask((worker) => {
                return new Promise((res, rej) => {
                    const channel = new MessageChannel();
                    channel.port1.onmessage = ({ data }) => {
                        channel.port1.close();
                        if (data.error) {
                            rej(data.error);
                        }
                        else {
                            res(data.result);
                        }
                    };
                    worker.postMessage({ operation, args, port: channel.port2 }, [channel.port2, ...(transferList ?? [])]);
                });
            });
        },
        close: () => {
            return Promise.all([
                ...workers.map(({ worker }) => {
                    return new Promise((res) => {
                        worker.addListener("exit", (code) => {
                            res(code);
                        });
                    });
                }),
                (async () => {
                    closed.abort();
                })(),
            ]);
        }
    };
};
export const withWorkerThreads = (taskCaller) => (...options) => async (fn) => {
    if (isMainThread) {
        const workerpool = makeWorkerPool(...options);
        try {
            const poolOps = Object.fromEntries(Object.entries(taskCaller).map(([k, v]) => [k, (...args) => {
                    return v(workerpool.task(k))(...args);
                }]));
            return await fn(poolOps);
        }
        finally {
            await workerpool.close();
        }
    }
    else {
        return await fn(undefined);
    }
};
export const implementWorker = (operations) => {
    if (!isMainThread) {
        parentPort.on("message", async ({ close, operation, args, port }) => {
            try {
                if (close) {
                    parentPort.close();
                }
                else {
                    const res = await operations[operation](...args);
                    if (typeof res === "object" && "result" in res) {
                        port.postMessage({ result: res.result }, res.transfer ?? []);
                    }
                    else {
                        port.postMessage({ result: res }, []);
                    }
                }
            }
            catch (e) {
                port?.postMessage({ error: e });
            }
        });
    }
};
//# sourceMappingURL=index.js.map