ci: shard extension fast checks

scripts/lib/extension-test-plan.mjs

@@ -2,8 +2,10 @@ import fs from "node:fs";
 import path from "node:path";
 import { channelTestRoots } from "../../vitest.channel-paths.mjs";
 import { BUNDLED_PLUGIN_PATH_PREFIX, BUNDLED_PLUGIN_ROOT_DIR } from "./bundled-plugin-paths.mjs";
+import { listAvailableExtensionIds } from "./changed-extensions.mjs";
 
 const repoRoot = path.resolve(import.meta.dirname, "..", "..");
+export const DEFAULT_EXTENSION_TEST_SHARD_COUNT = 6;
 
 function normalizeRelative(inputPath) {
   return inputPath.split(path.sep).join("/");
@@ -102,3 +104,100 @@ export function resolveExtensionTestPlan(params = {}) {
     testFileCount,
   };
 }
+
+function mergeTestPlans(plans) {
+  const groupsByConfig = new Map();
+
+  for (const plan of plans) {
+    const current = groupsByConfig.get(plan.config) ?? {
+      config: plan.config,
+      extensionIds: [],
+      roots: [],
+      testFileCount: 0,
+    };
+
+    current.extensionIds.push(plan.extensionId);
+    current.roots.push(...plan.roots);
+    current.testFileCount += plan.testFileCount;
+    groupsByConfig.set(plan.config, current);
+  }
+
+  const planGroups = [...groupsByConfig.values()]
+    .map((group) => ({
+      ...group,
+      extensionIds: group.extensionIds.toSorted((left, right) => left.localeCompare(right)),
+      roots: [...new Set(group.roots)],
+    }))
+    .toSorted((left, right) => left.config.localeCompare(right.config));
+
+  return {
+    extensionCount: plans.length,
+    extensionIds: plans
+      .map((plan) => plan.extensionId)
+      .toSorted((left, right) => left.localeCompare(right)),
+    hasTests: plans.length > 0,
+    planGroups,
+    testFileCount: plans.reduce((sum, plan) => sum + plan.testFileCount, 0),
+  };
+}
+
+export function resolveExtensionBatchPlan(params = {}) {
+  const cwd = params.cwd ?? process.cwd();
+  const extensionIds = params.extensionIds ?? listAvailableExtensionIds();
+  const plans = extensionIds
+    .map((extensionId) => resolveExtensionTestPlan({ cwd, targetArg: extensionId }))
+    .filter((plan) => plan.hasTests);
+
+  return mergeTestPlans(plans);
+}
+
+function pickLeastLoadedShard(shards) {
+  return shards.reduce((bestIndex, shard, index) => {
+    if (bestIndex === -1) {
+      return index;
+    }
+    const best = shards[bestIndex];
+    if (shard.testFileCount !== best.testFileCount) {
+      return shard.testFileCount < best.testFileCount ? index : bestIndex;
+    }
+    if (shard.plans.length !== best.plans.length) {
+      return shard.plans.length < best.plans.length ? index : bestIndex;
+    }
+    return index < bestIndex ? index : bestIndex;
+  }, -1);
+}
+
+export function createExtensionTestShards(params = {}) {
+  const cwd = params.cwd ?? process.cwd();
+  const extensionIds = params.extensionIds ?? listAvailableExtensionIds();
+  const shardCount = Math.max(1, Number.parseInt(String(params.shardCount ?? ""), 10) || 1);
+  const plans = extensionIds
+    .map((extensionId) => resolveExtensionTestPlan({ cwd, targetArg: extensionId }))
+    .filter((plan) => plan.hasTests)
+    .toSorted((left, right) => {
+      if (left.testFileCount !== right.testFileCount) {
+        return right.testFileCount - left.testFileCount;
+      }
+      return left.extensionId.localeCompare(right.extensionId);
+    });
+
+  const effectiveShardCount = Math.min(shardCount, Math.max(1, plans.length));
+  const shards = Array.from({ length: effectiveShardCount }, () => ({
+    plans: [],
+    testFileCount: 0,
+  }));
+
+  for (const plan of plans) {
+    const targetIndex = pickLeastLoadedShard(shards);
+    shards[targetIndex].plans.push(plan);
+    shards[targetIndex].testFileCount += plan.testFileCount;
+  }
+
+  return shards
+    .map((shard, index) => ({
+      index,
+      checkName: `checks-fast-extensions-shard-${index + 1}`,
+      ...mergeTestPlans(shard.plans),
+    }))
+    .filter((shard) => shard.hasTests);
+}

scripts/test-extension-batch.mjs

@@ -0,0 +1,105 @@
+#!/usr/bin/env node
+
+import { spawn } from "node:child_process";
+import path from "node:path";
+import { fileURLToPath, pathToFileURL } from "node:url";
+import { resolveExtensionBatchPlan } from "./lib/extension-test-plan.mjs";
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+const repoRoot = path.resolve(__dirname, "..");
+const pnpm = "pnpm";
+
+async function runVitestBatch(params) {
+  return await new Promise((resolve, reject) => {
+    const child = spawn(
+      pnpm,
+      ["exec", "vitest", "run", "--config", params.config, ...params.targets, ...params.args],
+      {
+        cwd: repoRoot,
+        stdio: "inherit",
+        shell: process.platform === "win32",
+        env: params.env,
+      },
+    );
+
+    child.on("error", reject);
+    child.on("exit", (code, signal) => {
+      if (signal) {
+        process.kill(process.pid, signal);
+        return;
+      }
+      resolve(code ?? 1);
+    });
+  });
+}
+
+function printUsage() {
+  console.error("Usage: pnpm test:extensions:batch <extension[,extension...]> [vitest args...]");
+  console.error(
+    "       node scripts/test-extension-batch.mjs <extension[,extension...]> [vitest args...]",
+  );
+}
+
+function parseExtensionIds(rawArgs) {
+  const args = [...rawArgs];
+  const extensionIds = [];
+
+  while (args[0] && !args[0].startsWith("-")) {
+    extensionIds.push(
+      ...args
+        .shift()
+        .split(",")
+        .map((value) => value.trim())
+        .filter(Boolean),
+    );
+  }
+
+  return { extensionIds, passthroughArgs: args };
+}
+
+async function run() {
+  const rawArgs = process.argv.slice(2);
+  if (rawArgs.includes("--help") || rawArgs.includes("-h")) {
+    printUsage();
+    return;
+  }
+
+  const passthroughArgs = rawArgs.filter((arg) => arg !== "--");
+  const { extensionIds, passthroughArgs: vitestArgs } = parseExtensionIds(passthroughArgs);
+  if (extensionIds.length === 0) {
+    printUsage();
+    process.exit(1);
+  }
+
+  const batchPlan = resolveExtensionBatchPlan({ cwd: process.cwd(), extensionIds });
+  if (!batchPlan.hasTests) {
+    console.log("[test-extension-batch] No tests found for the requested extensions. Skipping.");
+    return;
+  }
+
+  console.log(
+    `[test-extension-batch] Running ${batchPlan.testFileCount} test files across ${batchPlan.extensionCount} extensions`,
+  );
+
+  for (const group of batchPlan.planGroups) {
+    console.log(
+      `[test-extension-batch] ${group.config}: ${group.extensionIds.join(", ")} (${group.testFileCount} files)`,
+    );
+    const exitCode = await runVitestBatch({
+      args: vitestArgs,
+      config: group.config,
+      env: process.env,
+      targets: group.roots,
+    });
+    if (exitCode !== 0) {
+      process.exit(exitCode);
+    }
+  }
+}
+
+const entryHref = process.argv[1] ? pathToFileURL(path.resolve(process.argv[1])).href : "";
+
+if (import.meta.url === entryHref) {
+  await run();
+}