|
| 1 | +// Authored as plain CommonJS (NOT .ts) on purpose. vitest loads each package's vitest.config.ts by |
| 2 | +// bundling it, and it EXTERNALIZES this workspace subpath - node then loads this file verbatim. A .ts |
| 3 | +// here reaches node as raw TypeScript and crashes config loading on CI's pinned node 20 (no type |
| 4 | +// stripping: `SyntaxError`). Keeping it dependency-free JS - and importing nothing from the ESM-only |
| 5 | +// `vitest/node` - makes it loadable on every node. Types for consumers live in sequencer.d.cts. |
| 6 | + |
| 7 | +const { existsSync, readFileSync } = require("node:fs"); |
| 8 | +const path = require("node:path"); |
| 9 | + |
| 10 | +// Walk up from the package dir (cwd at config-load time) to the monorepo root (pnpm-workspace.yaml). |
| 11 | +function findRepoRoot(start) { |
| 12 | + let dir = start; |
| 13 | + for (let i = 0; i < 20; i++) { |
| 14 | + if (existsSync(path.join(dir, "pnpm-workspace.yaml"))) return dir; |
| 15 | + const parent = path.dirname(dir); |
| 16 | + if (parent === dir) break; |
| 17 | + dir = parent; |
| 18 | + } |
| 19 | + return start; |
| 20 | +} |
| 21 | + |
| 22 | +// test-timings.json lives at the monorepo root: { "<repo-relative path>": <ms> } |
| 23 | +const REPO_ROOT = findRepoRoot(process.cwd()); |
| 24 | +const TIMINGS_PATH = path.resolve(REPO_ROOT, "test-timings.json"); |
| 25 | + |
| 26 | +let cachedTimings; |
| 27 | + |
| 28 | +function loadTimings() { |
| 29 | + if (!cachedTimings) { |
| 30 | + cachedTimings = existsSync(TIMINGS_PATH) ? JSON.parse(readFileSync(TIMINGS_PATH, "utf-8")) : {}; |
| 31 | + } |
| 32 | + return cachedTimings; |
| 33 | +} |
| 34 | + |
| 35 | +function median(nums) { |
| 36 | + if (nums.length === 0) return 1; |
| 37 | + const sorted = [...nums].sort((a, b) => a - b); |
| 38 | + const mid = Math.floor(sorted.length / 2); |
| 39 | + return sorted.length % 2 ? sorted[mid] : (sorted[mid - 1] + sorted[mid]) / 2; |
| 40 | +} |
| 41 | + |
| 42 | +/** |
| 43 | + * Duration-weighted interpretation of `--shard=i/N`. Instead of vitest's default file-count split, |
| 44 | + * this greedily bin-packs test files by recorded duration (test-timings.json at the repo root; |
| 45 | + * unknown/new files get the median) so each shard does roughly equal work. |
| 46 | + * |
| 47 | + * The packing is fully deterministic (sort by duration desc, then moduleId) so every shard computes |
| 48 | + * the identical bins and just takes its own - no file runs twice or gets dropped. Falls back to the |
| 49 | + * full set when no shard is configured, and to ~count-based when no timings exist. |
| 50 | + * |
| 51 | + * Implemented as a standalone TestSequencer (not extending BaseSequencer) so this file never imports |
| 52 | + * `vitest/node` - see the header note. |
| 53 | + */ |
| 54 | +class DurationShardingSequencer { |
| 55 | + constructor(ctx) { |
| 56 | + this.ctx = ctx; |
| 57 | + } |
| 58 | + |
| 59 | + // Deterministic order (heaviest first, then moduleId) - stable across shards and a sensible |
| 60 | + // in-shard run order, replacing BaseSequencer's default sort we no longer inherit. |
| 61 | + async sort(files) { |
| 62 | + const timings = loadTimings(); |
| 63 | + const fallback = median(Object.values(timings)); |
| 64 | + return [...files].sort((a, b) => { |
| 65 | + const am = timings[path.relative(REPO_ROOT, a.moduleId)] ?? fallback; |
| 66 | + const bm = timings[path.relative(REPO_ROOT, b.moduleId)] ?? fallback; |
| 67 | + return bm - am || a.moduleId.localeCompare(b.moduleId); |
| 68 | + }); |
| 69 | + } |
| 70 | + |
| 71 | + async shard(specs) { |
| 72 | + const shard = this.ctx.config.shard; |
| 73 | + if (!shard || specs.length === 0) { |
| 74 | + return specs; |
| 75 | + } |
| 76 | + |
| 77 | + const timings = loadTimings(); |
| 78 | + const fallback = median(Object.values(timings)); |
| 79 | + |
| 80 | + const weighted = specs |
| 81 | + .map((spec) => ({ |
| 82 | + spec, |
| 83 | + ms: timings[path.relative(REPO_ROOT, spec.moduleId)] ?? fallback, |
| 84 | + })) |
| 85 | + .sort((a, b) => b.ms - a.ms || a.spec.moduleId.localeCompare(b.spec.moduleId)); |
| 86 | + |
| 87 | + const bins = Array.from({ length: shard.count }, () => ({ total: 0, specs: [] })); |
| 88 | + |
| 89 | + for (const { spec, ms } of weighted) { |
| 90 | + const lightest = bins.reduce((min, bin) => (bin.total < min.total ? bin : min)); |
| 91 | + lightest.total += ms; |
| 92 | + lightest.specs.push(spec); |
| 93 | + } |
| 94 | + |
| 95 | + return bins[shard.index - 1].specs; |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +module.exports = { DurationShardingSequencer }; |
0 commit comments