82 lines
2.3 KiB
TypeScript
82 lines
2.3 KiB
TypeScript
import { createRequire } from "node:module";
|
|
import os from "node:os";
|
|
import path from "node:path";
|
|
import type { IdDrawStyle } from "./id-config";
|
|
|
|
export type WorkerFont = {
|
|
path: string;
|
|
family: string;
|
|
};
|
|
|
|
export type GenerateTask = {
|
|
text: string;
|
|
outPath: string;
|
|
};
|
|
|
|
const nodeRequire = createRequire(path.join(process.cwd(), "package.json"));
|
|
|
|
function spawnWorker(workerPath: string, workerData: unknown) {
|
|
const threads = nodeRequire("node:worker_threads") as typeof import("node:worker_threads");
|
|
return new threads.Worker(workerPath, { workerData });
|
|
}
|
|
|
|
function workerCount(taskCount: number): number {
|
|
const cpus =
|
|
typeof os.availableParallelism === "function"
|
|
? os.availableParallelism()
|
|
: os.cpus().length;
|
|
return Math.max(1, Math.min(cpus, 8, taskCount));
|
|
}
|
|
|
|
function splitTasks<T>(items: T[], parts: number): T[][] {
|
|
const chunks: T[][] = Array.from({ length: parts }, () => []);
|
|
items.forEach((item, index) => {
|
|
chunks[index % parts]?.push(item);
|
|
});
|
|
return chunks.filter((chunk) => chunk.length > 0);
|
|
}
|
|
|
|
export async function generatePngsInWorkers(options: {
|
|
imagePath: string;
|
|
fonts: WorkerFont[];
|
|
style: IdDrawStyle;
|
|
tasks: GenerateTask[];
|
|
onProgress: (completed: number, total: number) => void;
|
|
}): Promise<void> {
|
|
const { imagePath, fonts, style, tasks, onProgress } = options;
|
|
if (tasks.length === 0) return;
|
|
|
|
const workerPath = path.join(process.cwd(), "workers", "generate-id.cjs");
|
|
const chunks = splitTasks(tasks, workerCount(tasks.length));
|
|
let completed = 0;
|
|
|
|
await Promise.all(
|
|
chunks.map(
|
|
(chunk) =>
|
|
new Promise<void>((resolve, reject) => {
|
|
const worker = spawnWorker(workerPath, {
|
|
imagePath,
|
|
fonts,
|
|
style,
|
|
tasks: chunk,
|
|
});
|
|
|
|
worker.on("message", (message: { type?: string }) => {
|
|
if (message?.type !== "progress") return;
|
|
completed += 1;
|
|
onProgress(completed, tasks.length);
|
|
});
|
|
|
|
worker.once("error", reject);
|
|
worker.once("exit", (code) => {
|
|
if (code === 0) {
|
|
resolve();
|
|
return;
|
|
}
|
|
reject(new Error(`Image worker stopped with code ${code}`));
|
|
});
|
|
}),
|
|
),
|
|
);
|
|
}
|