massive-cloud/server/algorithmJob.ts

/**
 * 通用算法任务调度器（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 } 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 false;
  }

  /** 取消某个 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 this.activeChildren) {
      const jobTaskId = (child as ChildProcess & { _taskId?: number })._taskId;
      if (jobTaskId === taskId) {
        child.kill('SIGKILL');
        cancelled.push(resultId);
      }
    }
    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);

/**
 * 把一次插件运行加入并发池。返回的 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);
}