297 lines
9.3 KiB
TypeScript
297 lines
9.3 KiB
TypeScript
/**
|
||
* 通用算法任务调度器（Route B）
|
||
* ------------------------------------------------------------
|
||
* 职责：
|
||
* 1. 维护一个带上限的并发池（默认 4 并发），避免一次性 fork 太多
|
||
* Python 子进程把宿主机打爆；
|
||
* 2. 把每个任务 fork 到独立 Node 子进程（algorithmJobRunner.ts），
|
||
* 让 HTTP/Vite 主进程不阻塞；
|
||
* 3. 子进程负责按 manifest 执行 process / mock runtime，
|
||
* 解析 stdout 进度，写 result.json，更新数据库。
|
||
*
|
||
* 主进程不直接 spawn Python，而是 fork Node runner —— 与现有
|
||
* smoothHeatmapJob / deviationHeatmapJob 的架构保持一致。
|
||
*/
|
||
import { fork, ChildProcess, execFileSync } from 'node:child_process';
|
||
import { fileURLToPath } from 'node:url';
|
||
import path from 'node:path';
|
||
|
||
const DEFAULT_CONCURRENCY = Number(process.env.ALGORITHM_JOB_CONCURRENCY ?? 4);
|
||
|
||
interface PendingJob {
|
||
resultId: number;
|
||
taskId: number;
|
||
pluginId: string;
|
||
params: Record<string, unknown>;
|
||
/** 可选：只处理选中的点云文件 ID；undefined/空数组 = 全部 */
|
||
pointCloudFileIds?: number[];
|
||
/** 可选：任务级 CAD 文件 ID；undefined/空数组 = 不传 CAD */
|
||
cadFileIds?: number[];
|
||
resolve: () => void;
|
||
reject: (err: unknown) => void;
|
||
}
|
||
|
||
class ConcurrencyPool {
|
||
private running = 0;
|
||
private queue: PendingJob[] = [];
|
||
private readonly limit: number;
|
||
private readonly activeChildren = new Map<number, ChildProcess>(); // resultId → child
|
||
|
||
constructor(limit: number) {
|
||
this.limit = Math.max(1, limit);
|
||
}
|
||
|
||
get stats() {
|
||
return { running: this.running, queued: this.queue.length, limit: this.limit };
|
||
}
|
||
|
||
enqueue(job: Omit<PendingJob, 'resolve' | 'reject'>): Promise<void> {
|
||
return new Promise((resolve, reject) => {
|
||
this.queue.push({ ...job, resolve, reject });
|
||
this.pump();
|
||
});
|
||
}
|
||
|
||
/** 取消指定 resultId 的 job：若在队列中直接移除；若已在运行则 kill 子进程。 */
|
||
cancel(resultId: number): boolean {
|
||
// 先从队列里找（还未开始执行）
|
||
const idx = this.queue.findIndex((j) => j.resultId === resultId);
|
||
if (idx !== -1) {
|
||
const [job] = this.queue.splice(idx, 1);
|
||
job.resolve(); // 让等待方正常退出
|
||
return true;
|
||
}
|
||
// 再找正在运行的子进程
|
||
const child = this.activeChildren.get(resultId);
|
||
if (child) {
|
||
child.kill('SIGKILL');
|
||
return true;
|
||
}
|
||
return killOrphanAlgorithmProcesses({ resultId });
|
||
}
|
||
|
||
/** 取消某个 taskId 下所有 job（删除任务时用）。返回被取消的 resultId 数组。 */
|
||
cancelByTask(taskId: number): number[] {
|
||
const cancelled: number[] = [];
|
||
// 队列中移除
|
||
const keep: PendingJob[] = [];
|
||
for (const job of this.queue) {
|
||
if (job.taskId === taskId) {
|
||
job.resolve();
|
||
cancelled.push(job.resultId);
|
||
} else {
|
||
keep.push(job);
|
||
}
|
||
}
|
||
this.queue = keep;
|
||
// 运行中 kill
|
||
for (const [resultId, child] of Array.from(this.activeChildren.entries())) {
|
||
const jobTaskId = (child as ChildProcess & { _taskId?: number })._taskId;
|
||
if (jobTaskId === taskId) {
|
||
child.kill('SIGKILL');
|
||
cancelled.push(resultId);
|
||
}
|
||
}
|
||
const orphaned = killOrphanAlgorithmProcesses({ taskId });
|
||
if (orphaned.length > 0) {
|
||
cancelled.push(...orphaned);
|
||
}
|
||
return cancelled;
|
||
}
|
||
|
||
private pump(): void {
|
||
while (this.running < this.limit && this.queue.length > 0) {
|
||
const next = this.queue.shift()!;
|
||
this.running += 1;
|
||
this.spawnChild(next).finally(() => {
|
||
this.running = Math.max(0, this.running - 1);
|
||
this.pump();
|
||
});
|
||
}
|
||
}
|
||
|
||
private async spawnChild(job: PendingJob): Promise<void> {
|
||
const script = resolveAlgorithmRunnerEntry();
|
||
const isProd = process.env.NODE_ENV === 'production';
|
||
const args = [
|
||
String(job.resultId),
|
||
String(job.taskId),
|
||
job.pluginId,
|
||
JSON.stringify(job.params ?? {}),
|
||
JSON.stringify(job.pointCloudFileIds ?? []),
|
||
JSON.stringify(job.cadFileIds ?? []),
|
||
];
|
||
const child = fork(script, args, {
|
||
env: { ...process.env },
|
||
stdio: ['ignore', 'inherit', 'inherit', 'ipc'],
|
||
execArgv: isProd ? [] : ['--import', 'tsx'],
|
||
});
|
||
// 把 taskId 附到 child 上，cancelByTask 用
|
||
(child as ChildProcess & { _taskId?: number })._taskId = job.taskId;
|
||
this.activeChildren.set(job.resultId, child);
|
||
|
||
return new Promise<void>((resolve) => {
|
||
child.on('error', (err) => {
|
||
console.error('[AlgorithmJob] fork 失败:', err);
|
||
});
|
||
child.on('exit', (code, signal) => {
|
||
this.activeChildren.delete(job.resultId);
|
||
if (code !== 0 && code != null) {
|
||
console.error(
|
||
`[AlgorithmJob] 子进程退出 code=${code} signal=${signal ?? ''} pluginId=${job.pluginId} resultId=${job.resultId}`,
|
||
);
|
||
}
|
||
// 即使失败也 resolve：runner 已经把数据库改成 error，主进程无需再处理
|
||
job.resolve();
|
||
resolve();
|
||
});
|
||
});
|
||
}
|
||
}
|
||
|
||
function resolveAlgorithmRunnerEntry(): string {
|
||
const dir = path.dirname(fileURLToPath(import.meta.url));
|
||
if (process.env.NODE_ENV === 'production') {
|
||
return path.join(dir, 'algorithmJobRunner.js');
|
||
}
|
||
return path.join(dir, 'algorithmJobRunner.ts');
|
||
}
|
||
|
||
const pool = new ConcurrencyPool(DEFAULT_CONCURRENCY);
|
||
|
||
interface ProcessRow {
|
||
pid: number;
|
||
ppid: number;
|
||
command: string;
|
||
}
|
||
|
||
function listProcesses(): ProcessRow[] {
|
||
try {
|
||
const output = execFileSync('ps', ['-eo', 'pid=,ppid=,command='], {
|
||
encoding: 'utf8',
|
||
maxBuffer: 1024 * 1024,
|
||
});
|
||
return output
|
||
.split('\n')
|
||
.map((line) => line.trim())
|
||
.filter(Boolean)
|
||
.map((line) => {
|
||
const match = line.match(/^(\d+)\s+(\d+)\s+(.+)$/);
|
||
if (!match) return null;
|
||
return {
|
||
pid: Number(match[1]),
|
||
ppid: Number(match[2]),
|
||
command: match[3],
|
||
};
|
||
})
|
||
.filter((row): row is ProcessRow => row != null && Number.isFinite(row.pid) && Number.isFinite(row.ppid));
|
||
} catch (err) {
|
||
console.warn('[AlgorithmJob] 无法列出进程，跳过孤儿进程清理:', err);
|
||
return [];
|
||
}
|
||
}
|
||
|
||
function collectDescendants(processes: ProcessRow[], rootPids: Set<number>): Set<number> {
|
||
const all = new Set(rootPids);
|
||
let changed = true;
|
||
while (changed) {
|
||
changed = false;
|
||
for (const proc of processes) {
|
||
if (!all.has(proc.pid) && all.has(proc.ppid)) {
|
||
all.add(proc.pid);
|
||
changed = true;
|
||
}
|
||
}
|
||
}
|
||
return all;
|
||
}
|
||
|
||
function killPids(pids: number[]): number[] {
|
||
const killed: number[] = [];
|
||
for (const pid of pids) {
|
||
if (pid <= 1 || pid === process.pid) continue;
|
||
try {
|
||
process.kill(pid, 'SIGKILL');
|
||
killed.push(pid);
|
||
} catch {
|
||
// Process may already have exited.
|
||
}
|
||
}
|
||
return killed;
|
||
}
|
||
|
||
function killOrphanAlgorithmProcesses(input: { resultId: number }): boolean;
|
||
function killOrphanAlgorithmProcesses(input: { taskId: number }): number[];
|
||
function killOrphanAlgorithmProcesses(input: { resultId?: number; taskId?: number }): boolean | number[] {
|
||
const processes = listProcesses();
|
||
const roots = new Set<number>();
|
||
const resultIds = new Set<number>();
|
||
|
||
for (const proc of processes) {
|
||
const cmd = proc.command;
|
||
if (input.resultId != null) {
|
||
const resultId = String(input.resultId);
|
||
if (
|
||
new RegExp(`algorithmJobRunner\\.(?:ts|js)\\s+${resultId}\\s+`).test(cmd) ||
|
||
cmd.includes(`/analysis/${resultId}/input.json`)
|
||
) {
|
||
roots.add(proc.pid);
|
||
resultIds.add(input.resultId);
|
||
}
|
||
continue;
|
||
}
|
||
|
||
if (input.taskId != null) {
|
||
const runnerMatch = cmd.match(/algorithmJobRunner\.(?:ts|js)\s+(\d+)\s+(\d+)\s+/);
|
||
if (runnerMatch && Number(runnerMatch[2]) === input.taskId) {
|
||
roots.add(proc.pid);
|
||
resultIds.add(Number(runnerMatch[1]));
|
||
continue;
|
||
}
|
||
const inputJsonMatch = cmd.match(new RegExp(`/data/tasks/${input.taskId}/analysis/(\\d+)/input\\.json`));
|
||
if (inputJsonMatch) {
|
||
roots.add(proc.pid);
|
||
resultIds.add(Number(inputJsonMatch[1]));
|
||
}
|
||
}
|
||
}
|
||
|
||
const targets = Array.from(collectDescendants(processes, roots)).sort((a, b) => b - a);
|
||
const killed = killPids(targets);
|
||
if (killed.length > 0) {
|
||
console.warn(`[AlgorithmJob] 已清理孤儿算法进程 pid=${killed.join(', ')}`);
|
||
}
|
||
|
||
if (input.resultId != null) return killed.length > 0;
|
||
return Array.from(resultIds);
|
||
}
|
||
|
||
/**
|
||
* 把一次插件运行加入并发池。返回的 Promise 在子进程退出时 resolve（无论成功/失败）。
|
||
* 真实结果状态写在 task_analysis_results.resultData / status 中，由前端轮询消费。
|
||
*/
|
||
export function enqueueAlgorithmJob(
|
||
resultId: number,
|
||
taskId: number,
|
||
pluginId: string,
|
||
params: Record<string, unknown> = {},
|
||
pointCloudFileIds?: number[],
|
||
cadFileIds?: number[],
|
||
): Promise<void> {
|
||
return pool.enqueue({ resultId, taskId, pluginId, params, pointCloudFileIds, cadFileIds });
|
||
}
|
||
|
||
export function algorithmJobStats() {
|
||
return pool.stats;
|
||
}
|
||
|
||
/** 取消单个分析结果对应的 job（删除分析结果时调用）。 */
|
||
export function cancelAlgorithmJob(resultId: number): boolean {
|
||
return pool.cancel(resultId);
|
||
}
|
||
|
||
/** 取消某检测任务下所有 job（删除任务时调用）。 */
|
||
export function cancelAlgorithmJobsByTask(taskId: number): number[] {
|
||
return pool.cancelByTask(taskId);
|
||
}
|