mirror of
https://github.com/mattermost/mattermost.git
synced 2026-05-24 03:05:38 -04:00
294 lines
12 KiB
JavaScript
294 lines
12 KiB
JavaScript
#!/usr/bin/env node
|
|
/**
|
|
* shard-split.js — Test shard assignment solver
|
|
*
|
|
* Splits Go test packages across N parallel CI runners using timing data
|
|
* from previous runs. Uses a two-tier strategy:
|
|
*
|
|
* 1. "Light" packages (< HEAVY_MS total runtime): assigned whole to a shard
|
|
* 2. "Heavy" packages (>= HEAVY_MS): individual tests distributed across
|
|
* shards using -run regex filters
|
|
*
|
|
* Timing data sources (in priority order):
|
|
* - gotestsum.json (JSONL): per-test elapsed times from previous run
|
|
* - prev-report.xml (JUnit XML): package-level timing (fallback)
|
|
* - Round-robin: when no timing data exists at all
|
|
*
|
|
* Assignment algorithm: greedy bin-packing (sort by duration desc, assign
|
|
* each item to the shard with lowest current load). Simple and effective
|
|
* for our distribution where 2 packages dominate 84% of runtime.
|
|
*
|
|
* Environment variables:
|
|
* SHARD_INDEX — this runner's index (0-based)
|
|
* SHARD_TOTAL — total number of shards
|
|
*
|
|
* Input files (in working directory):
|
|
* all-packages.txt — newline-separated list of all test packages
|
|
* prev-gotestsum.json — (optional) JSONL timing data from previous run
|
|
* prev-report.xml — (optional) JUnit XML from previous run
|
|
*
|
|
* Output files (in working directory):
|
|
* shard-te-packages.txt — space-separated TE packages for this shard
|
|
* shard-ee-packages.txt — space-separated EE packages for this shard
|
|
* shard-heavy-runs.txt — heavy package runs, one per line: "pkg REGEX"
|
|
*/
|
|
|
|
const fs = require("node:fs");
|
|
const { execSync } = require("node:child_process");
|
|
|
|
const SHARD_INDEX = parseInt(process.env.SHARD_INDEX);
|
|
const SHARD_TOTAL = parseInt(process.env.SHARD_TOTAL);
|
|
const HEAVY_MS = 600000; // 600s (10 min): packages above this get test-level splitting
|
|
|
|
// Packages that should always be split test-by-test, even on a cold cache.
|
|
// Without timing data the splitter falls through to alphabetical round-robin,
|
|
// which places these adjacent on the same runner and overwhelms postgres.
|
|
// Forcing them heavy lets `go test -list` enumerate their tests so the
|
|
// bin-packer can spread them across all shards.
|
|
const KNOWN_HEAVY_PKGS = new Set([
|
|
"github.com/mattermost/mattermost/server/v8/channels/api4",
|
|
"github.com/mattermost/mattermost/server/v8/channels/app",
|
|
]);
|
|
|
|
if (isNaN(SHARD_INDEX) || isNaN(SHARD_TOTAL) || SHARD_TOTAL < 1) {
|
|
console.error("ERROR: SHARD_INDEX and SHARD_TOTAL must be set");
|
|
process.exit(1);
|
|
}
|
|
|
|
const allPkgs = fs
|
|
.readFileSync("all-packages.txt", "utf8")
|
|
.trim()
|
|
.split("\n")
|
|
.filter(Boolean);
|
|
if (allPkgs.length === 0) {
|
|
console.error("WARNING: No test packages found in all-packages.txt");
|
|
process.exit(0);
|
|
}
|
|
|
|
const pkgTimes = {};
|
|
const testTimes = {}; // "pkg::TestName" -> ms
|
|
|
|
// ── Parse gotestsum.json (JSONL) for per-test timing ──
|
|
// Each line is a JSON event; we want "pass" events with Elapsed times.
|
|
if (fs.existsSync("prev-gotestsum.json")) {
|
|
console.log("::group::Parsing gotestsum.json timing data");
|
|
const lines = fs.readFileSync("prev-gotestsum.json", "utf8").split("\n");
|
|
for (const line of lines) {
|
|
if (!line.includes('"pass"')) continue;
|
|
try {
|
|
const d = JSON.parse(line);
|
|
if (!d.Test || !d.Package) continue;
|
|
const elapsed = Math.round((d.Elapsed || 0) * 1000);
|
|
// Aggregate package time from test pass events
|
|
pkgTimes[d.Package] = (pkgTimes[d.Package] || 0) + elapsed;
|
|
// Top-level test name (use max elapsed for parent vs subtests)
|
|
const top = d.Test.split("/")[0];
|
|
const key = d.Package + "::" + top;
|
|
testTimes[key] = Math.max(testTimes[key] || 0, elapsed);
|
|
} catch (e) {
|
|
// Skip malformed lines
|
|
}
|
|
}
|
|
console.log(
|
|
`gotestsum.json: ${Object.keys(pkgTimes).length} packages, ${Object.keys(testTimes).length} tests`,
|
|
);
|
|
console.log("::endgroup::");
|
|
}
|
|
|
|
// ── Fallback: parse JUnit XML for package-level timing ──
|
|
if (Object.keys(pkgTimes).length === 0 && fs.existsSync("prev-report.xml")) {
|
|
console.log("::group::Parsing JUnit XML timing data (fallback)");
|
|
const xml = fs.readFileSync("prev-report.xml", "utf8");
|
|
for (const m of xml.matchAll(/<testsuite[^>]*>/g)) {
|
|
const name = m[0].match(/name="([^"]+)"/)?.[1];
|
|
const time = m[0].match(/\btime="([^"]+)"/)?.[1];
|
|
if (name && time) {
|
|
pkgTimes[name] =
|
|
(pkgTimes[name] || 0) + Math.round(parseFloat(time) * 1000);
|
|
}
|
|
}
|
|
console.log(
|
|
`JUnit XML: ${Object.keys(pkgTimes).length} packages (no per-test data)`,
|
|
);
|
|
console.log("::endgroup::");
|
|
}
|
|
|
|
const hasTimingData = Object.keys(pkgTimes).length > 0;
|
|
const hasTestTiming = Object.keys(testTimes).length > 0;
|
|
|
|
// ── Identify heavy packages ──
|
|
// Split at test level for packages above HEAVY_MS (requires per-test timing)
|
|
// AND for the KNOWN_HEAVY_PKGS list (which uses go test -list discovery
|
|
// to enumerate tests when no timing cache exists).
|
|
//
|
|
// Both checks gate on allPkgs membership so stale entries from the cached
|
|
// pkgTimes (renamed/deleted packages from a prior run) can't end up in
|
|
// heavyPkgs — otherwise the post-discovery fallback would emit them as
|
|
// whole-package items for nonexistent packages.
|
|
const allPkgsSet = new Set(allPkgs);
|
|
const heavyPkgs = new Set();
|
|
if (hasTestTiming) {
|
|
for (const [pkg, ms] of Object.entries(pkgTimes)) {
|
|
if (ms > HEAVY_MS && allPkgsSet.has(pkg)) heavyPkgs.add(pkg);
|
|
}
|
|
}
|
|
for (const pkg of allPkgs) {
|
|
if (KNOWN_HEAVY_PKGS.has(pkg)) heavyPkgs.add(pkg);
|
|
}
|
|
if (heavyPkgs.size > 0) {
|
|
console.log("Heavy packages (test-level splitting):");
|
|
for (const p of heavyPkgs) {
|
|
const t = pkgTimes[p];
|
|
const label = t ? `${(t / 1000).toFixed(0)}s` : "no-timing";
|
|
console.log(` ${label} ${p.split("/").pop()}`);
|
|
}
|
|
}
|
|
|
|
// ── Build work items ──
|
|
// Each item is either a whole package ("P") or a single test from a heavy package ("T")
|
|
const items = [];
|
|
for (const pkg of allPkgs) {
|
|
if (heavyPkgs.has(pkg)) {
|
|
// Split into individual test items
|
|
const tests = Object.entries(testTimes)
|
|
.filter(([k]) => k.startsWith(pkg + "::"))
|
|
.map(([k, ms]) => ({ ms, type: "T", pkg, test: k.split("::")[1] }));
|
|
if (tests.length > 0) {
|
|
items.push(...tests);
|
|
}
|
|
// If no per-test timing exists, the discovery step below enumerates
|
|
// tests via `go test -list`. A final fallback to whole-package is
|
|
// added after discovery for packages where both lookups failed.
|
|
} else {
|
|
items.push({ ms: pkgTimes[pkg] || 1, type: "P", pkg });
|
|
}
|
|
}
|
|
// ── Discover new/renamed tests in heavy packages ──
|
|
// Tests not in the timing cache won't appear in any shard's -run regex,
|
|
// silently skipping them. Discover current test names at runtime and
|
|
// assign any cache-missing tests to the least-loaded shard.
|
|
if (heavyPkgs.size > 0) {
|
|
console.log("::group::Discovering new tests in heavy packages");
|
|
for (const pkg of heavyPkgs) {
|
|
const cachedTests = new Set(
|
|
Object.keys(testTimes)
|
|
.filter((k) => k.startsWith(pkg + "::"))
|
|
.map((k) => k.split("::")[1]),
|
|
);
|
|
try {
|
|
const out = execSync(`go test -list '.*' ${pkg} 2>/dev/null`, {
|
|
encoding: "utf8",
|
|
timeout: 300000,
|
|
});
|
|
const currentTests = out
|
|
.split("\n")
|
|
.map((l) => l.trim())
|
|
.filter((l) => /^Test[A-Z]/.test(l));
|
|
let newCount = 0;
|
|
for (const t of currentTests) {
|
|
if (!cachedTests.has(t)) {
|
|
// Assign a small default duration so it gets picked up
|
|
items.push({ ms: 1000, type: "T", pkg, test: t });
|
|
newCount++;
|
|
}
|
|
}
|
|
if (newCount > 0) {
|
|
console.log(
|
|
` ${pkg.split("/").pop()}: ${newCount} new test(s) not in cache`,
|
|
);
|
|
}
|
|
} catch (e) {
|
|
// go test -list can fail for packages whose TestMain requires a DB
|
|
// connection (e.g. sqlstore) because the GitHub runner cannot reach the
|
|
// docker-compose postgres network. Log a warning and fall back to
|
|
// treating the package as a whole unit rather than failing all shards.
|
|
console.error(
|
|
`::warning::${pkg.split("/").pop()}: go test -list failed — treating as whole package (new tests may be skipped this run). ${e.message}`,
|
|
);
|
|
}
|
|
}
|
|
// Ensure every heavy package has at least one item. A package can reach
|
|
// this point with zero items if it has no per-test timing AND `go test
|
|
// -list` failed (e.g. sqlstore on a cold cache).
|
|
for (const pkg of heavyPkgs) {
|
|
const hasItems = items.some((it) => it.pkg === pkg);
|
|
if (!hasItems) {
|
|
console.log(
|
|
` ${pkg.split("/").pop()}: no per-test data, running as whole package`,
|
|
);
|
|
items.push({ ms: pkgTimes[pkg] || 1, type: "P", pkg });
|
|
}
|
|
}
|
|
console.log("::endgroup::");
|
|
}
|
|
|
|
// Sort descending by duration for greedy bin-packing
|
|
items.sort((a, b) => b.ms - a.ms);
|
|
|
|
// ── Greedy bin-packing assignment ──
|
|
const shards = Array.from({ length: SHARD_TOTAL }, () => ({
|
|
load: 0,
|
|
whole: [],
|
|
heavy: {},
|
|
}));
|
|
|
|
if (!hasTimingData && heavyPkgs.size === 0) {
|
|
// Round-robin fallback only when we have *no* signal — no timing cache
|
|
// and no known-heavy packages to test-level-split. With heavyPkgs we
|
|
// can still bin-pack: discovered tests (ms=1000 each) drive the
|
|
// distribution and whole-package items (ms=1) fill in evenly.
|
|
console.log("No timing data — using round-robin");
|
|
allPkgs.forEach((pkg, i) => {
|
|
shards[i % SHARD_TOTAL].whole.push(pkg);
|
|
});
|
|
} else {
|
|
for (const item of items) {
|
|
// Find shard with minimum current load
|
|
const min = shards.reduce(
|
|
(m, s, i) => (s.load < shards[m].load ? i : m),
|
|
0,
|
|
);
|
|
shards[min].load += item.ms;
|
|
if (item.type === "P") {
|
|
shards[min].whole.push(item.pkg);
|
|
} else {
|
|
if (!shards[min].heavy[item.pkg]) shards[min].heavy[item.pkg] = [];
|
|
shards[min].heavy[item.pkg].push(item.test);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ── Report shard assignments ──
|
|
console.log("::group::Shard assignment");
|
|
for (let i = 0; i < SHARD_TOTAL; i++) {
|
|
const s = shards[i];
|
|
const hRuns = Object.keys(s.heavy).length;
|
|
const hTests = Object.values(s.heavy).reduce((n, a) => n + a.length, 0);
|
|
const marker = i === SHARD_INDEX ? " ← THIS SHARD" : "";
|
|
console.log(
|
|
`Shard ${i}: ${(s.load / 1000).toFixed(1)}s | ${s.whole.length} pkgs` +
|
|
(hRuns > 0 ? `, ${hRuns} heavy splits (${hTests} tests)` : "") +
|
|
marker,
|
|
);
|
|
}
|
|
console.log("::endgroup::");
|
|
|
|
// ── Write output for this shard ──
|
|
const myShard = shards[SHARD_INDEX];
|
|
const te = myShard.whole.filter((p) => !p.includes("/enterprise/")).join(" ");
|
|
const ee = myShard.whole.filter((p) => p.includes("/enterprise/")).join(" ");
|
|
|
|
fs.writeFileSync("shard-te-packages.txt", te);
|
|
fs.writeFileSync("shard-ee-packages.txt", ee);
|
|
|
|
// Heavy package runs: one line per run as "pkg REGEX"
|
|
const heavyRuns = Object.entries(myShard.heavy).map(([pkg, tests]) => {
|
|
const regex = tests.map((t) => "^" + t + "$").join("|");
|
|
return pkg + " " + regex;
|
|
});
|
|
fs.writeFileSync("shard-heavy-runs.txt", heavyRuns.join("\n"));
|
|
|
|
console.log(
|
|
`Light packages: ${myShard.whole.length} (${te.split(" ").filter(Boolean).length} TE, ${ee.split(" ").filter(Boolean).length} EE)`,
|
|
);
|
|
console.log(`Heavy package runs: ${heavyRuns.length}`);
|