perf: speed up test parallelism

2026-04-24 07:01:49 +00:00 · 2026-03-26 20:08:49 +00:00
parent 2fc017788c
commit 663ba5a3cd
13 changed files with 1517 additions and 99 deletions
--- a/scripts/test-planner/executor.mjs
+++ b/scripts/test-planner/executor.mjs
@@ -132,6 +132,62 @@ export function formatExplanation(explanation) {
  ].join("\n");
 }

+const buildOrderedParallelSegments = (units) => {
+  const segments = [];
+  let deferredUnits = [];
+  for (const unit of units) {
+    if (unit.serialPhase) {
+      if (deferredUnits.length > 0) {
+        segments.push({ type: "deferred", units: deferredUnits });
+        deferredUnits = [];
+      }
+      const lastSegment = segments.at(-1);
+      if (lastSegment?.type === "serialPhase" && lastSegment.phase === unit.serialPhase) {
+        lastSegment.units.push(unit);
+      } else {
+        segments.push({ type: "serialPhase", phase: unit.serialPhase, units: [unit] });
+      }
+      continue;
+    }
+    deferredUnits.push(unit);
+  }
+  if (deferredUnits.length > 0) {
+    segments.push({ type: "deferred", units: deferredUnits });
+  }
+  return segments;
+};
+
+const prioritizeDeferredUnitsForPhase = (units, phase) => {
+  const preferredSurface =
+    phase === "extensions" || phase === "channels" ? phase : phase === "unit-fast" ? "unit" : null;
+  if (preferredSurface === null) {
+    return units;
+  }
+  const preferred = [];
+  const remaining = [];
+  for (const unit of units) {
+    if (unit.surface === preferredSurface) {
+      preferred.push(unit);
+    } else {
+      remaining.push(unit);
+    }
+  }
+  return preferred.length > 0 ? [...preferred, ...remaining] : units;
+};
+
+const partitionUnitsBySurface = (units, surface) => {
+  const matching = [];
+  const remaining = [];
+  for (const unit of units) {
+    if (unit.surface === surface) {
+      matching.push(unit);
+    } else {
+      remaining.push(unit);
+    }
+  }
+  return { matching, remaining };
+};
+
 export async function executePlan(plan, options = {}) {
  const env = options.env ?? process.env;
  const artifacts = options.artifacts ?? createExecutionArtifacts(env);
@@ -632,23 +688,116 @@ export async function executePlan(plan, options = {}) {
  }

  if (plan.serialPrefixUnits.length > 0) {
-    const failedSerialPrefix = await runUnitsWithLimit(
-      plan.serialPrefixUnits,
-      plan.passthroughOptionArgs,
-      1,
-    );
-    if (failedSerialPrefix !== undefined) {
-      return failedSerialPrefix;
+    const orderedSegments = buildOrderedParallelSegments(plan.parallelUnits);
+    let pendingDeferredSegment = null;
+    let carriedDeferredPromise = null;
+    let carriedDeferredSurface = null;
+    for (const segment of orderedSegments) {
+      if (segment.type === "deferred") {
+        pendingDeferredSegment = segment;
+        continue;
+      }
+      // Preserve phase ordering, but let batches inside the same shared phase use
+      // the normal top-level concurrency budget.
+      let deferredPromise = null;
+      let deferredCarryPromise = carriedDeferredPromise;
+      let deferredCarrySurface = carriedDeferredSurface;
+      if (
+        segment.phase === "unit-fast" &&
+        pendingDeferredSegment !== null &&
+        plan.topLevelParallelEnabled
+      ) {
+        const availableSlots = Math.max(0, plan.topLevelParallelLimit - segment.units.length);
+        if (availableSlots > 0) {
+          const prePhaseDeferred = pendingDeferredSegment.units;
+          if (prePhaseDeferred.length > 0) {
+            deferredCarryPromise = runUnitsWithLimit(
+              prePhaseDeferred,
+              plan.passthroughOptionArgs,
+              availableSlots,
+            );
+            deferredCarrySurface = prePhaseDeferred.some((unit) => unit.surface === "channels")
+              ? "channels"
+              : null;
+            pendingDeferredSegment = null;
+          }
+        }
+      }
+      if (pendingDeferredSegment !== null) {
+        const prioritizedDeferred = prioritizeDeferredUnitsForPhase(
+          pendingDeferredSegment.units,
+          segment.phase,
+        );
+        if (segment.phase === "extensions") {
+          const { matching: channelDeferred, remaining: otherDeferred } = partitionUnitsBySurface(
+            prioritizedDeferred,
+            "channels",
+          );
+          deferredPromise =
+            otherDeferred.length > 0
+              ? runUnitsWithLimit(
+                  otherDeferred,
+                  plan.passthroughOptionArgs,
+                  plan.deferredRunConcurrency ?? 1,
+                )
+              : null;
+          deferredCarryPromise =
+            channelDeferred.length > 0
+              ? runUnitsWithLimit(
+                  channelDeferred,
+                  plan.passthroughOptionArgs,
+                  plan.deferredRunConcurrency ?? 1,
+                )
+              : carriedDeferredPromise;
+          deferredCarrySurface = channelDeferred.length > 0 ? "channels" : carriedDeferredSurface;
+        } else {
+          deferredPromise = runUnitsWithLimit(
+            prioritizedDeferred,
+            plan.passthroughOptionArgs,
+            plan.deferredRunConcurrency ?? 1,
+          );
+        }
+      }
+      pendingDeferredSegment = null;
+      // eslint-disable-next-line no-await-in-loop
+      const failedSerialPhase = await runUnits(segment.units, plan.passthroughOptionArgs);
+      if (failedSerialPhase !== undefined) {
+        return failedSerialPhase;
+      }
+      if (deferredCarryPromise !== null && deferredCarrySurface === segment.phase) {
+        // eslint-disable-next-line no-await-in-loop
+        const failedCarriedDeferred = await deferredCarryPromise;
+        if (failedCarriedDeferred !== undefined) {
+          return failedCarriedDeferred;
+        }
+        deferredCarryPromise = null;
+        deferredCarrySurface = null;
+      }
+      if (deferredPromise !== null) {
+        // eslint-disable-next-line no-await-in-loop
+        const failedDeferredPhase = await deferredPromise;
+        if (failedDeferredPhase !== undefined) {
+          return failedDeferredPhase;
+        }
+      }
+      carriedDeferredPromise = deferredCarryPromise;
+      carriedDeferredSurface = deferredCarrySurface;
    }
-    const failedDeferredParallel = plan.deferredRunConcurrency
-      ? await runUnitsWithLimit(
-          plan.deferredParallelUnits,
-          plan.passthroughOptionArgs,
-          plan.deferredRunConcurrency,
-        )
-      : await runUnits(plan.deferredParallelUnits, plan.passthroughOptionArgs);
-    if (failedDeferredParallel !== undefined) {
-      return failedDeferredParallel;
+    if (pendingDeferredSegment !== null) {
+      const failedDeferredParallel = await runUnitsWithLimit(
+        pendingDeferredSegment.units,
+        plan.passthroughOptionArgs,
+        plan.deferredRunConcurrency ?? 1,
+      );
+      if (failedDeferredParallel !== undefined) {
+        return failedDeferredParallel;
+      }
+    }
+    if (carriedDeferredPromise !== null) {
+      const failedCarriedDeferred = await carriedDeferredPromise;
+      if (failedCarriedDeferred !== undefined) {
+        return failedCarriedDeferred;
+      }
    }
  } else {
    const failedParallel = await runUnits(plan.parallelUnits, plan.passthroughOptionArgs);
--- a/scripts/test-planner/planner.mjs
+++ b/scripts/test-planner/planner.mjs
@@ -2,6 +2,7 @@ import path from "node:path";
 import { isUnitConfigTestFile } from "../../vitest.unit-paths.mjs";
 import {
  loadChannelTimingManifest,
+  loadExtensionTimingManifest,
  loadUnitMemoryHotspotManifest,
  loadUnitTimingManifest,
  packFilesByDuration,
@@ -145,6 +146,7 @@ const createPlannerContext = (request, options = {}) => {
  const catalog = options.catalog ?? loadTestCatalog();
  const unitTimingManifest = loadUnitTimingManifest();
  const channelTimingManifest = loadChannelTimingManifest();
+  const extensionTimingManifest = loadExtensionTimingManifest();
  const unitMemoryHotspotManifest = loadUnitMemoryHotspotManifest();
  return {
    env,
@@ -153,6 +155,7 @@ const createPlannerContext = (request, options = {}) => {
    catalog,
    unitTimingManifest,
    channelTimingManifest,
+    extensionTimingManifest,
    unitMemoryHotspotManifest,
  };
 };
@@ -198,11 +201,16 @@ const resolveEntryTimingEstimator = (entry, context) => {
      context.unitTimingManifest.files[file]?.durationMs ??
      context.unitTimingManifest.defaultDurationMs;
  }
-  if (config === "vitest.channels.config.ts" || config === "vitest.extensions.config.ts") {
+  if (config === "vitest.channels.config.ts") {
    return (file) =>
      context.channelTimingManifest.files[file]?.durationMs ??
      context.channelTimingManifest.defaultDurationMs;
  }
+  if (config === "vitest.extensions.config.ts") {
+    return (file) =>
+      context.extensionTimingManifest.files[file]?.durationMs ??
+      context.extensionTimingManifest.defaultDurationMs;
+  }
  return null;
 };

@@ -233,6 +241,20 @@ const splitFilesByDurationBudget = (files, targetDurationMs, estimateDurationMs)
  return batches;
 };

+const splitFilesByBalancedDurationBudget = (files, targetDurationMs, estimateDurationMs) => {
+  if (!Number.isFinite(targetDurationMs) || targetDurationMs <= 0 || files.length <= 1) {
+    return [files];
+  }
+  const totalDurationMs = files.reduce((sum, file) => sum + estimateDurationMs(file), 0);
+  const batchCount = clamp(Math.ceil(totalDurationMs / targetDurationMs), 1, files.length);
+  const originalOrder = new Map(files.map((file, index) => [file, index]));
+  return packFilesByDuration(files, batchCount, estimateDurationMs).map((batch) =>
+    [...batch].toSorted(
+      (left, right) => (originalOrder.get(left) ?? 0) - (originalOrder.get(right) ?? 0),
+    ),
+  );
+};
+
 const resolveMaxWorkersForUnit = (unit, context) => {
  const overrideWorkers = Number.parseInt(context.env.OPENCLAW_TEST_WORKERS ?? "", 10);
  const resolvedOverride =
@@ -332,11 +354,20 @@ const resolveUnitHeavyFileGroups = (context) => {
 };

 const buildDefaultUnits = (context, request) => {
-  const { env, executionBudget, catalog, unitTimingManifest, channelTimingManifest } = context;
+  const {
+    env,
+    executionBudget,
+    catalog,
+    unitTimingManifest,
+    channelTimingManifest,
+    extensionTimingManifest,
+  } = context;
  const noIsolateArgs = context.noIsolateArgs;
  const selectedSurfaces = buildRequestedSurfaces(request, env);
  const selectedSurfaceSet = new Set(selectedSurfaces);
+  const unitOnlyRun = selectedSurfaceSet.size === 1 && selectedSurfaceSet.has("unit");
  const channelsOnlyRun = selectedSurfaceSet.size === 1 && selectedSurfaceSet.has("channels");
+  const extensionsOnlyRun = selectedSurfaceSet.size === 1 && selectedSurfaceSet.has("extensions");

  const {
    heavyUnitLaneCount,
@@ -361,6 +392,8 @@ const buildDefaultUnits = (context, request) => {
    unitTimingManifest.files[file]?.durationMs ?? unitTimingManifest.defaultDurationMs;
  const estimateChannelDurationMs = (file) =>
    channelTimingManifest.files[file]?.durationMs ?? channelTimingManifest.defaultDurationMs;
+  const estimateExtensionDurationMs = (file) =>
+    extensionTimingManifest.files[file]?.durationMs ?? extensionTimingManifest.defaultDurationMs;
  const unitFastCandidateFiles = catalog.allKnownUnitFiles.filter(
    (file) => !new Set(unitFastExcludedFiles).has(file),
  );
@@ -421,7 +454,7 @@ const buildDefaultUnits = (context, request) => {
            id: unitId,
            surface: "unit",
            isolate: false,
-            serialPhase: "unit-fast",
+            serialPhase: unitOnlyRun ? undefined : "unit-fast",
            includeFiles: batch,
            estimatedDurationMs: estimateEntryFilesDurationMs(
              { args: ["vitest", "run", "--config", "vitest.unit.config.ts"] },
@@ -453,6 +486,7 @@ const buildDefaultUnits = (context, request) => {
          id: `unit-${path.basename(file, ".test.ts")}-isolated`,
          surface: "unit",
          isolate: true,
+          estimatedDurationMs: estimateUnitDurationMs(file),
          args: [
            "vitest",
            "run",
@@ -478,6 +512,7 @@ const buildDefaultUnits = (context, request) => {
          id: `unit-heavy-${String(index + 1)}`,
          surface: "unit",
          isolate: false,
+          estimatedDurationMs: files.reduce((sum, file) => sum + estimateUnitDurationMs(file), 0),
          args: [
            "vitest",
            "run",
@@ -498,6 +533,7 @@ const buildDefaultUnits = (context, request) => {
          id: `unit-${path.basename(file, ".test.ts")}-memory-isolated`,
          surface: "unit",
          isolate: true,
+          estimatedDurationMs: estimateUnitDurationMs(file),
          args: [
            "vitest",
            "run",
@@ -533,6 +569,29 @@ const buildDefaultUnits = (context, request) => {
    }
  }

+  if (selectedSurfaceSet.has("channels")) {
+    for (const file of catalog.channelIsolatedFiles) {
+      units.push(
+        createExecutionUnit(context, {
+          id: `${path.basename(file, ".test.ts")}-channels-isolated`,
+          surface: "channels",
+          isolate: true,
+          estimatedDurationMs: estimateChannelDurationMs(file),
+          args: [
+            "vitest",
+            "run",
+            "--config",
+            "vitest.channels.config.ts",
+            "--pool=forks",
+            ...noIsolateArgs,
+            file,
+          ],
+          reasons: ["channels-isolated-rule"],
+        }),
+      );
+    }
+  }
+
  if (selectedSurfaceSet.has("extensions")) {
    for (const file of catalog.extensionForkIsolatedFiles) {
      units.push(
@@ -540,15 +599,16 @@ const buildDefaultUnits = (context, request) => {
          id: `extensions-${path.basename(file, ".test.ts")}-isolated`,
          surface: "extensions",
          isolate: true,
+          estimatedDurationMs: estimateExtensionDurationMs(file),
          args: ["vitest", "run", "--config", "vitest.extensions.config.ts", "--pool=forks", file],
          reasons: ["extensions-isolated-manifest"],
        }),
      );
    }
-    const extensionBatches = splitFilesByDurationBudget(
+    const extensionBatches = splitFilesByBalancedDurationBudget(
      extensionSharedCandidateFiles,
      extensionsBatchTargetMs,
-      estimateChannelDurationMs,
+      estimateExtensionDurationMs,
    );
    for (const [batchIndex, batch] of extensionBatches.entries()) {
      if (batch.length === 0) {
@@ -561,7 +621,7 @@ const buildDefaultUnits = (context, request) => {
          id: unitId,
          surface: "extensions",
          isolate: false,
-          serialPhase: "extensions",
+          serialPhase: extensionsOnlyRun ? undefined : "extensions",
          includeFiles: batch,
          estimatedDurationMs: estimateEntryFilesDurationMs(
            { args: ["vitest", "run", "--config", "vitest.extensions.config.ts"] },
@@ -581,25 +641,6 @@ const buildDefaultUnits = (context, request) => {
  }

  if (selectedSurfaceSet.has("channels")) {
-    for (const file of catalog.channelIsolatedFiles) {
-      units.push(
-        createExecutionUnit(context, {
-          id: `${path.basename(file, ".test.ts")}-channels-isolated`,
-          surface: "channels",
-          isolate: true,
-          args: [
-            "vitest",
-            "run",
-            "--config",
-            "vitest.channels.config.ts",
-            "--pool=forks",
-            ...noIsolateArgs,
-            file,
-          ],
-          reasons: ["channels-isolated-rule"],
-        }),
-      );
-    }
    const channelBatches = splitFilesByDurationBudget(
      channelSharedCandidateFiles,
      channelsBatchTargetMs,
--- a/scripts/test-planner/runtime-profile.mjs
+++ b/scripts/test-planner/runtime-profile.mjs
@@ -102,6 +102,7 @@ const LOCAL_MEMORY_BUDGETS = {
    heavyLaneCount: 3,
    memoryHeavyFileLimit: 8,
    unitFastBatchTargetMs: 10_000,
+    channelsBatchTargetMs: 0,
  },
  moderate: {
    vitestCap: 3,
@@ -117,6 +118,7 @@ const LOCAL_MEMORY_BUDGETS = {
    heavyLaneCount: 4,
    memoryHeavyFileLimit: 12,
    unitFastBatchTargetMs: 15_000,
+    channelsBatchTargetMs: 0,
  },
  mid: {
    vitestCap: 4,
@@ -132,6 +134,7 @@ const LOCAL_MEMORY_BUDGETS = {
    heavyLaneCount: 4,
    memoryHeavyFileLimit: 16,
    unitFastBatchTargetMs: 0,
+    channelsBatchTargetMs: 0,
  },
  high: {
    vitestCap: 6,
@@ -140,13 +143,14 @@ const LOCAL_MEMORY_BUDGETS = {
    unitHeavy: 2,
    extensions: 4,
    gateway: 3,
-    topLevelNoIsolate: 12,
+    topLevelNoIsolate: 14,
    topLevelIsolated: 4,
-    deferred: 3,
+    deferred: 8,
    heavyFileLimit: 80,
    heavyLaneCount: 5,
    memoryHeavyFileLimit: 16,
    unitFastBatchTargetMs: 45_000,
+    channelsBatchTargetMs: 30_000,
  },
 };

@@ -296,8 +300,8 @@ export function resolveExecutionBudget(runtimeCapabilities) {
    memoryHeavyUnitFileLimit: bandBudget.memoryHeavyFileLimit,
    unitFastLaneCount: 1,
    unitFastBatchTargetMs: bandBudget.unitFastBatchTargetMs,
-    channelsBatchTargetMs: 0,
-    extensionsBatchTargetMs: 0,
+    channelsBatchTargetMs: bandBudget.channelsBatchTargetMs ?? 0,
+    extensionsBatchTargetMs: 240_000,
  };

  const loadAdjustedBudget = {