Genarrative/apps/ai-game-creator-shell/scripts/run-rust-shell-test-shards.mjs

#!/usr/bin/env node
// AGC 壳 Rust 套件的分片运行器。
//
// 背景：AGC 壳的 Rust 单测集中在 src/main.rs 的 bin target（实测 2466 条），仓库口径用
// `--test-threads=1` 跑，理由是并行调度会让共享 Agent Runtime 后台锁与异步终态的用例在
// **同一进程内**互相干扰（见 development-workflow 的 Tauri suite 单线程口径）。代价是
// 整套用例串行跑满 507 秒，占掉 CI 上 `AI game creator shell Rust tests` job 的大头。
//
// 这里保留「片内串行」的既有口径，只把用例集合切成 N 片：
// - CI 用 `--shard-index=<i>` 让**每个 job 只跑一片**，靠多个 job 并发把整套用例摊开；
// - 本地不传 `--shard-index` 时把 N 片放进 N 个**独立进程**并行（--concurrency 可调），
//   保留一条命令跑全量的入口。
// 片并集必须等于全集、且不得重复，数量不符即失败，防止分片规则改动后静默漏跑；该校验
// 与「只跑一片」无关，因此在每个 job 上都会执行。
//
// 注意：同一容器内多进程并行这套用例（共享 HOME、target、固定临时路径）实测会互相拖慢，
// 比串行还慢，所以 CI 走「一个 job 一片」而不是单 job 内并行。
//
// 用法：
//   node scripts/run-rust-shell-test-shards.mjs --shards=4 --shard-index=2   # CI：只跑第 2 片
//   node scripts/run-rust-shell-test-shards.mjs [--shards=4] [--concurrency=4]  # 本地：全量
//   node scripts/run-rust-shell-test-shards.mjs --manifest=<Cargo.toml> --target-kind=lib --no-locked
//
// 参数：
//   --shards=<n>        分片数，默认 4
//   --shard-index=<i>   只跑第 i 片（1..shards）；不传则跑全部分片
//   --concurrency=<n>   同时运行的片数，默认等于分片数；--shard-index 时恒为 1
//   --manifest=<path>   Cargo.toml，默认 ../src-tauri/Cargo.toml（相对本脚本）
//   --target-kind=<k>   bin | lib，默认 bin（本地自测小 crate 时用 lib）
//   --bin=<name>        bin target 名，默认 genarrative-ai-game-creator-shell
//   --package=<name>    target-kind=lib 时要跑的包名（配合 lib 目标使用）
//   --no-locked         传给 cargo 时不带 --locked（只对没有提交 Cargo.lock 的 crate 需要）
//   --shard-tmp-root=<path>  片专属 TMPDIR 的父目录，默认 <系统临时目录>/agc-rust-shards

import { spawn } from 'node:child_process';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { fileURLToPath } from 'node:url';

const scriptDirectory = path.dirname(fileURLToPath(import.meta.url));
const shellRoot = path.resolve(scriptDirectory, '..');

function fail(message) {
  console.error(`[rust-shards] ${message}`);
  process.exit(1);
}

function parsePositiveInteger(name, rawValue) {
  if (!/^[1-9][0-9]*$/.test(rawValue)) {
    fail(`${name} must be a positive integer, received: ${rawValue}`);
  }

  return Number(rawValue);
}

const options = {
  shards: 4,
  shardIndex: undefined,
  concurrency: undefined,
  manifestPath: path.join(shellRoot, 'src-tauri', 'Cargo.toml'),
  targetKind: 'bin',
  binName: 'genarrative-ai-game-creator-shell',
  packageName: undefined,
  locked: true,
  shardTmpRoot: path.join(os.tmpdir(), 'agc-rust-shards'),
};

for (const rawArgument of process.argv.slice(2)) {
  if (rawArgument === '--no-locked') {
    options.locked = false;
    continue;
  }

  const separatorIndex = rawArgument.indexOf('=');
  if (!rawArgument.startsWith('--') || separatorIndex === -1) {
    fail(`unexpected argument: ${rawArgument}`);
  }

  const name = rawArgument.slice(2, separatorIndex);
  const value = rawArgument.slice(separatorIndex + 1);
  switch (name) {
    case 'shards':
      options.shards = parsePositiveInteger('--shards', value);
      break;
    case 'shard-index':
      options.shardIndex = parsePositiveInteger('--shard-index', value);
      break;
    case 'concurrency':
      options.concurrency = parsePositiveInteger('--concurrency', value);
      break;
    case 'manifest':
      options.manifestPath = path.resolve(process.cwd(), value);
      break;
    case 'target-kind':
      if (value !== 'bin' && value !== 'lib') {
        fail(`--target-kind must be bin or lib, received: ${value}`);
      }
      options.targetKind = value;
      break;
    case 'bin':
      options.binName = value;
      break;
    case 'package':
      options.packageName = value;
      break;
    case 'shard-tmp-root':
      options.shardTmpRoot = path.resolve(process.cwd(), value);
      break;
    default:
      fail(`unexpected argument: ${rawArgument}`);
  }
}

if (!fs.existsSync(options.manifestPath)) {
  fail(`manifest does not exist: ${options.manifestPath}`);
}

if (options.shardIndex !== undefined && options.shardIndex > options.shards) {
  fail(
    `--shard-index (${options.shardIndex}) must be within --shards (${options.shards})`,
  );
}

const concurrency = options.concurrency ?? options.shards;
const crateRoot = path.dirname(options.manifestPath);

function buildCargoArguments(target) {
  const cargoArguments = ['test'];
  if (options.locked) {
    cargoArguments.push('--locked');
  }
  cargoArguments.push('--manifest-path', options.manifestPath);
  if (target.kind === 'lib') {
    if (options.packageName !== undefined) {
      cargoArguments.push('-p', options.packageName);
    }
    cargoArguments.push('--lib');
    return cargoArguments;
  }
  cargoArguments.push('--bin', target.name);
  return cargoArguments;
}

function formatDuration(milliseconds) {
  return `${(milliseconds / 1000).toFixed(1)}s`;
}

// 编译一次，直接拿到测试可执行文件：后续每片都运行同一个二进制，不再各自调用 cargo，
// 免得 N 个 cargo 去争 package cache 与 target 目录锁。
function resolveTestExecutable() {
  return new Promise((resolve, reject) => {
    const cargoArguments = buildCargoArguments({
      kind: options.targetKind,
      name: options.binName,
    });
    cargoArguments.push('--no-run', '--message-format=json');
    console.log(`[rust-shards] cargo ${cargoArguments.join(' ')}`);

    const child = spawn('cargo', cargoArguments, {
      cwd: crateRoot,
      env: process.env,
      stdio: ['ignore', 'pipe', 'inherit'],
    });

    let buffered = '';
    const executables = [];
    child.stdout.setEncoding('utf8');
    child.stdout.on('data', (chunk) => {
      buffered += chunk;
      const lines = buffered.split('\n');
      buffered = lines.pop() ?? '';
      for (const line of lines) {
        if (!line.startsWith('{')) {
          continue;
        }
        let message;
        try {
          message = JSON.parse(line);
        } catch {
          continue;
        }
        if (
          message.reason === 'compiler-artifact' &&
          typeof message.executable === 'string'
        ) {
          executables.push(message.executable);
        }
      }
    });

    child.on('error', (error) => {
      reject(new Error(`unable to start cargo: ${error.message}`));
    });
    child.on('close', (code) => {
      if (code !== 0) {
        reject(
          new Error(
            `cargo ${cargoArguments.join(' ')} failed with exit code ${code}`,
          ),
        );
        return;
      }
      const uniqueExecutables = [...new Set(executables)];
      if (uniqueExecutables.length !== 1) {
        reject(
          new Error(
            `expected exactly one test executable for the ${options.targetKind} target, found ${uniqueExecutables.length}: ${uniqueExecutables.join(', ')}`,
          ),
        );
        return;
      }
      resolve(uniqueExecutables[0]);
    });
  });
}

function listTestNames(executable) {
  return new Promise((resolve, reject) => {
    const child = spawn(executable, ['--list'], {
      cwd: crateRoot,
      env: process.env,
      stdio: ['ignore', 'pipe', 'pipe'],
    });

    let stdout = '';
    let stderr = '';
    child.stdout.setEncoding('utf8');
    child.stderr.setEncoding('utf8');
    child.stdout.on('data', (chunk) => {
      stdout += chunk;
    });
    child.stderr.on('data', (chunk) => {
      stderr += chunk;
    });
    child.on('error', (error) => {
      reject(new Error(`unable to list tests: ${error.message}`));
    });
    child.on('close', (code) => {
      if (code !== 0) {
        reject(
          new Error(
            `listing tests failed with exit code ${code}: ${stderr.trim()}`,
          ),
        );
        return;
      }
      const names = [];
      for (const line of stdout.split('\n')) {
        const match = /^(.*): test$/.exec(line.trim());
        if (match !== null && match[1].length > 0) {
          names.push(match[1]);
        }
      }
      resolve(names);
    });
  });
}

function splitTestNames(testNames, shardCount) {
  const sortedNames = [...testNames].sort();
  const shards = Array.from({ length: shardCount }, () => []);
  sortedNames.forEach((testName, index) => {
    shards[index % shardCount].push(testName);
  });
  return shards;
}

function assertShardsCoverEveryTest(testNames, shards) {
  const flattened = shards.flat();
  if (flattened.length !== testNames.length) {
    fail(
      `shard split covered ${flattened.length} of ${testNames.length} tests; the split rule must be exhaustive`,
    );
  }
  if (new Set(flattened).size !== flattened.length) {
    fail(
      'shard split selected the same test more than once; the split rule must be disjoint',
    );
  }
  const missing = testNames.filter((testName) => !flattened.includes(testName));
  if (missing.length > 0) {
    fail(
      `shard split is missing tests, for example: ${missing.slice(0, 5).join(', ')}`,
    );
  }
}

function runShard(executable, shardIndex, shardCount, shardTestNames) {
  const label = `shard ${shardIndex + 1}/${shardCount}`;
  const shardTmpDirectory = path.join(
    options.shardTmpRoot,
    `shard-${shardIndex + 1}`,
  );
  fs.rmSync(shardTmpDirectory, { recursive: true, force: true });
  fs.mkdirSync(shardTmpDirectory, { recursive: true });
  const startedAt = Date.now();

  return new Promise((resolve) => {
    const child = spawn(
      executable,
      ['--exact', ...shardTestNames, '--test-threads=1'],
      {
        cwd: crateRoot,
        env: {
          ...process.env,
          TMPDIR: shardTmpDirectory,
          TMP: shardTmpDirectory,
          TEMP: shardTmpDirectory,
        },
        stdio: ['ignore', 'pipe', 'pipe'],
      },
    );

    const failureLines = [];
    let inFailureList = false;
    let stderr = '';
    const consumeLine = (rawLine) => {
      const line = rawLine.replace(/\r$/, '');
      if (line.includes('failures:')) {
        inFailureList = true;
        return;
      }
      if (inFailureList) {
        if (line.trim().length === 0) {
          inFailureList = false;
          return;
        }
        failureLines.push(line.trim());
      }
    };

    child.stdout.setEncoding('utf8');
    child.stderr.setEncoding('utf8');
    let stdoutBuffer = '';
    child.stdout.on('data', (chunk) => {
      stdoutBuffer += chunk;
      const lines = stdoutBuffer.split('\n');
      stdoutBuffer = lines.pop() ?? '';
      for (const line of lines) {
        consumeLine(line);
      }
    });
    child.stderr.on('data', (chunk) => {
      stderr += chunk;
    });

    child.on('error', (error) => {
      resolve({
        label,
        ok: false,
        durationMs: Date.now() - startedAt,
        testCount: shardTestNames.length,
        failures: [`unable to start test binary: ${error.message}`],
        stderr,
      });
    });
    child.on('close', (code) => {
      resolve({
        label,
        ok: code === 0,
        durationMs: Date.now() - startedAt,
        testCount: shardTestNames.length,
        failures: failureLines,
        stderr,
      });
    });
  });
}

async function runWithConcurrency(shards, runner) {
  const results = new Array(shards.length);
  let nextIndex = 0;
  const workers = Array.from(
    { length: Math.min(concurrency, shards.length) },
    async () => {
      while (nextIndex < shards.length) {
        const index = nextIndex;
        nextIndex += 1;
        results[index] = await runner(shards[index], index);
      }
    },
  );
  await Promise.all(workers);
  return results;
}

async function main() {
  const executable = await resolveTestExecutable();
  const testNames = await listTestNames(executable);
  if (testNames.length === 0) {
    fail(
      `no tests discovered in ${options.manifestPath} (${options.targetKind})`,
    );
  }

  const shards = splitTestNames(testNames, options.shards);
  assertShardsCoverEveryTest(testNames, shards);

  const selectedShards =
    options.shardIndex === undefined
      ? shards.map((shardTestNames, index) => ({ index, shardTestNames }))
      : [
          {
            index: options.shardIndex - 1,
            shardTestNames: shards[options.shardIndex - 1],
          },
        ];

  if (options.shardIndex === undefined) {
    console.log(
      `[rust-shards] ${testNames.length} tests, ${shards.length} shard(s), concurrency ${Math.min(concurrency, shards.length)}`,
    );
    for (const { index, shardTestNames } of selectedShards) {
      console.log(
        `[rust-shards] shard ${index + 1}/${shards.length}: ${shardTestNames.length} test(s)`,
      );
    }
  } else {
    console.log(
      `[rust-shards] ${testNames.length} tests, ${shards.length} shard(s), running shard ${options.shardIndex}/${shards.length} (${selectedShards[0].shardTestNames.length} test(s))`,
    );
  }

  const results = await runWithConcurrency(
    selectedShards,
    ({ index, shardTestNames }) =>
      runShard(executable, index, shards.length, shardTestNames),
  );

  let failed = false;
  for (const result of results) {
    if (result.ok) {
      console.log(
        `[rust-shards] ${result.label} ok: ${result.testCount} test(s) in ${formatDuration(result.durationMs)}`,
      );
      continue;
    }
    failed = true;
    console.error(
      `[rust-shards] ${result.label} FAILED: ${result.testCount} test(s) in ${formatDuration(result.durationMs)}`,
    );
    for (const failure of result.failures) {
      console.error(`[rust-shards]   ${failure}`);
    }
    if (result.stderr.trim().length > 0) {
      console.error(
        `[rust-shards]   stderr: ${result.stderr.trim().split('\n').slice(-20).join('\n[rust-shards]   ')}`,
      );
    }
  }

  if (failed) {
    process.exit(1);
  }
  console.log('[rust-shards] OK');
}

main().catch((error) => {
  fail(error instanceof Error ? error.message : String(error));
});