ci: wire up integration-test sharding, off by default

.github/workflows/integration_tests.yml

@@ -14,8 +14,20 @@ on:
         default: true
     secrets:
       CODECOV_TOKEN: { required: false }
+env:
+  # Number of parallel shards each compression config is split into. Every
+  # top-level test lands in exactly one shard, so the shards together are the
+  # full suite — keep this in sync with strategy.matrix.shard below.
+  #
+  # Deliberately 1: this wires the sharding up without spending anything. The
+  # infra-tests pool is saturated, and measured at 4 shards the extra jobs
+  # spent their savings queueing for runners. Raise it (and the shard list) to
+  # spend new capacity once #3352 lands.
+  SHARDS: "1"
+
 jobs:
   run:
+    name: run (${{ matrix.name }}, shard ${{ matrix.shard }})
     if: ${{ inputs.run-tests == true }}
     runs-on: infra-tests
     timeout-minutes: 30
@@ -25,6 +37,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        name: [uncompressed, zstd1, lz4]
+        shard: [1]
         include:
           - name: uncompressed
             compress_enabled: "false"
@@ -86,6 +100,7 @@ jobs:
           TESTS_ORCHESTRATOR_HOST: "localhost:5008"
           TESTS_ENVD_PROXY: "http://localhost:3002"
           TESTS_CLIENT_PROXY: "http://localhost:3002"
+          SHARD: ${{ matrix.shard }}
         run: |
           # Run the integration tests
           make test-integration
@@ -123,7 +138,7 @@ jobs:
         if: ${{ always() && inputs.publish == true }}
         uses: actions/upload-artifact@v6
         with:
-          name: Integration Tests Results (${{ matrix.name }})
+          name: Integration Tests Results (${{ matrix.name }} ${{ matrix.shard }})
           path: ./tests/integration/test-results.xml
 
       - name: Upload test results to Codecov
@@ -140,7 +155,7 @@ jobs:
         if: ${{ always() && inputs.publish == true }}
         uses: actions/upload-artifact@v6
         with:
-          name: Service Logs (${{ matrix.name }})
+          name: Service Logs (${{ matrix.name }} ${{ matrix.shard }})
           path: ~/logs/*.log
 
   integration_tests:

tests/integration/Makefile

@@ -21,6 +21,15 @@ seed:
 		go run seed.go
 	@echo "Done"
 
+# Tunables, overridable from the environment (CI sets SHARDS/SHARD).
+# SHARDS > 1 runs only shard SHARD of the selected packages; the shards
+# partition the suite, so the union of all SHARDS runs is the full suite.
+TEST_PARALLEL := $(or $(strip $(TEST_PARALLEL)),4)
+TEST_TIMEOUT := $(or $(strip $(TEST_TIMEOUT)),20m)
+JUNIT_FILE := $(or $(strip $(JUNIT_FILE)),test-results.xml)
+SHARDS := $(or $(strip $(SHARDS)),1)
+SHARD := $(or $(strip $(SHARD)),1)
+
 .PHONY: test
 test: test/.
 test/%:
@@ -35,14 +44,24 @@ test/%:
 	export TESTS_SANDBOX_USER_ID=$(TESTS_SANDBOX_USER_ID); \
 	go test -v ./internal/main_test.go -count=1 && \
 	TEST_PATH="./internal/tests/$(subst test/,,$@)"; \
+	TEST_PATH="$${TEST_PATH%/.}"; \
+	RUN=""; \
 	case "$${TEST_PATH}" in \
-		*.go:*) \
-			BASE=$${TEST_PATH%%:*}; \
-			TEST_FN=$${TEST_PATH#*:}; \
-			go tool gotestsum --rerun-fails=1 --packages="$$BASE" --format standard-verbose --junitfile=test-results.xml -- -count=1 -parallel=4 -timeout=20m -run "$${TEST_FN}" ;; \
-		*.go) go tool gotestsum --rerun-fails=1 --packages="$$TEST_PATH" --format standard-verbose --junitfile=test-results.xml -- -count=1 -parallel=4 -timeout=20m ;; \
-		*) go tool gotestsum --rerun-fails=1 --packages="$$TEST_PATH/..." --format standard-verbose --junitfile=test-results.xml -- -count=1 -parallel=4 -timeout=20m ;; \
-	esac
+		*.go:*) PACKAGES="$${TEST_PATH%%:*}"; RUN="$${TEST_PATH#*:}" ;; \
+		*.go) PACKAGES="$${TEST_PATH}" ;; \
+		*) PACKAGES="$${TEST_PATH}/..." ;; \
+	esac; \
+	if [ "$(SHARDS)" -gt 1 ]; then RUN=$$(./scripts/select-tests.sh $(SHARD) $(SHARDS) "$${PACKAGES}") || exit 1; fi; \
+	if [ -n "$${RUN}" ]; then set -- -run "$${RUN}"; else set --; fi; \
+	go tool gotestsum --rerun-fails=1 --packages="$${PACKAGES}" --format standard-verbose --junitfile=$(JUNIT_FILE) \
+		-- -count=1 -parallel=$(TEST_PARALLEL) -timeout=$(TEST_TIMEOUT) "$$@"
+
+# Refresh the shard-balancing weights from the JUnit files of a full CI run:
+#   gh run download <run-id> -p 'Integration Tests Results*' -D /tmp/junit
+#   make -C tests/integration update-test-weights JUNIT_DIR=/tmp/junit
+.PHONY: update-test-weights
+update-test-weights:
+	go run scripts/update-test-weights.go $(JUNIT_DIR)
 
 .PHONY: connect-orchestrator
 connect-orchestrator:

tests/integration/README.md

@@ -9,6 +9,23 @@ Package for defining integration tests. Currently, there is a setup for API and
 3. If necessary, run `make connect-orchestrator` to create a tunnel to one orchestrator client VM in GCP (you may need to run `make setup-ssh` the first time)
 4. Run `make test` in this folder or `make test-integration` from the root `infra/` folder.
 
+Narrow the run with `make test/<path under internal/tests>`, e.g. `make test/api/templates`
+or `make test/api/templates/build_template_test.go:TestTemplateBuildCOPY`.
+
+## Sharding
+
+A compression config can be split across parallel CI jobs, each running one
+shard: `SHARDS=2 SHARD=1 make test`. `scripts/select-tests.sh` enumerates the
+top-level tests from the source tree and bin-packs each package across the
+shards using the recorded per-test times in `scripts/test-weights.tsv`, so every
+test runs in exactly one shard and the shards finish at roughly the same time.
+Enumeration happens at run time, so a newly added test always lands in exactly
+one shard; a test missing from the weights file still runs, it just gets a
+median-time estimate for balancing. Refresh the weights with
+`make update-test-weights JUNIT_DIR=…` when the suite's shape changes.
+
+CI ships `SHARDS: 1`, i.e. sharding off — see the note in the workflow.
+
 ## Usage of clients (api, orchestrator, envd)
 
 All tests are in the folder internal/tests. You can see the usage of different clients in the tests. Here are just basics.

tests/integration/scripts/select-tests.sh

@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+# Print the `go test -run` regex selecting shard $1 of $2 for the given packages.
+#
+#   select-tests.sh <shard> <shard-count> <package>...
+#
+# Every top-level test is enumerated from the source tree at run time, so a
+# newly added test always lands in exactly one shard: the shards partition the
+# suite, they never filter it. test-weights.tsv only decides *which* shard a
+# test goes to; a missing entry costs balance, never coverage.
+#
+# Balancing is per package: each package's tests are bin-packed (longest
+# processing time first) across all shards, so every shard gets a slice of
+# every package and keeps the cross-package parallelism `go test` gives us.
+set -euo pipefail
+
+if [ "$#" -lt 3 ]; then
+	echo "usage: $0 <shard-index-1-based> <shard-count> <package>..." >&2
+	exit 2
+fi
+
+shard=$1
+shards=$2
+shift 2
+
+script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+weights=${script_dir}/test-weights.tsv
+
+if [ "$shard" -lt 1 ] || [ "$shard" -gt "$shards" ]; then
+	echo "shard $shard out of range 1..$shards" >&2
+	exit 2
+fi
+
+dirs=()
+for pkg in "$@"; do
+	pkg=${pkg#./}
+	case "$pkg" in
+	*/...)
+		while IFS= read -r dir; do
+			dirs+=("${dir#./}")
+		done < <(find "${pkg%/...}" -type d | sort)
+		;;
+	*) dirs+=("$pkg") ;;
+	esac
+done
+
+# Median of the known weights, used for tests that predate the last refresh of
+# test-weights.tsv (or were added since).
+default_weight=$(awk -F'\t' '!/^#/ && NF == 3 { print $3 }' "$weights" | sort -n |
+	awk '{ v[NR] = $1 } END { if (NR) print v[int((NR + 1) / 2)]; else print 30 }')
+
+names=$(
+	for dir in "${dirs[@]}"; do
+		compgen -G "${dir}/*_test.go" >/dev/null || continue
+		sed -n 's/^func \(Test[A-Za-z0-9_]*\)(.*/\1/p' "${dir}"/*_test.go |
+			sort -u | sed "s|^|${dir}\t|"
+	done
+)
+
+if [ -z "$names" ]; then
+	echo "no tests found in: $*" >&2
+	exit 1
+fi
+
+
+# Attach a weight to every test, then hand the list to the packer sorted by
+# package, weight descending, name — the input order LPT bin-packing needs.
+printf '%s\n' "$names" |
+	awk -F'\t' -v weights="$weights" -v def="$default_weight" '
+		BEGIN { while ((getline line < weights) > 0) { n = split(line, f, "\t"); if (line !~ /^#/ && n == 3) w[f[1] SUBSEP f[2]] = f[3] } }
+		{ key = $1 SUBSEP $2; print $1 "\t" ((key in w) ? w[key] : def) "\t" $2 }
+	' |
+	sort -t"$(printf '\t')" -k1,1 -k2,2nr -k3,3 |
+	awk -F'\t' -v shard="$shard" -v shards="$shards" '
+		function reset(i) { for (i = 1; i <= shards; i++) load[i] = 0 }
+		BEGIN { reset(); pkg = "" }
+		{
+			if ($1 != pkg) { pkg = $1; reset() }
+			best = 1
+			for (i = 2; i <= shards; i++) if (load[i] < load[best]) best = i
+			load[best] += $2
+			if (best == shard) out = out (out ? "|" : "") $3
+		}
+		END {
+			if (out == "") { print "shard " shard " of " shards " matched no tests" > "/dev/stderr"; exit 1 }
+			print "^(" out ")$"
+		}
+	'