diff --git a/rust/private/rust.bzl b/rust/private/rust.bzl index 0dc2d993e7..dfc65fabbf 100644 --- a/rust/private/rust.bzl +++ b/rust/private/rust.bzl @@ -951,15 +951,6 @@ _RUST_TEST_ATTRS = { "env_inherit": attr.string_list( doc = "Specifies additional environment variables to inherit from the external environment when the test is executed by bazel test.", ), - "use_libtest_harness": attr.bool( - mandatory = False, - default = True, - doc = dedent("""\ - Whether to use `libtest`. For targets using this flag, individual tests can be run by using the - [--test_arg](https://docs.bazel.build/versions/4.0.0/command-line-reference.html#flag--test_arg) flag. - E.g. `bazel test //src:rust_test --test_arg=foo::test::test_fn`. - """), - ), "experimental_enable_sharding": attr.bool( mandatory = False, default = False, @@ -968,14 +959,25 @@ _RUST_TEST_ATTRS = { When enabled, tests are executed via a wrapper script that: 1. Enumerates tests using libtest's --list flag - 2. Partitions tests across shards based on TEST_SHARD_INDEX/TEST_TOTAL_SHARDS - 3. Runs only the tests assigned to the current shard + 2. Sorts tests by name and partitions them across shards by stable name hash + 3. Uses either Bazel's native TEST_TOTAL_SHARDS/TEST_SHARD_INDEX env + or explicit RULES_RUST_TEST_TOTAL_SHARDS/RULES_RUST_TEST_SHARD_INDEX env + 4. Runs only the tests assigned to the current shard This attribute only has an effect when use_libtest_harness is True. This is experimental and may change in future releases. """), ), + "use_libtest_harness": attr.bool( + mandatory = False, + default = True, + doc = dedent("""\ + Whether to use `libtest`. For targets using this flag, individual tests can be run by using the + [--test_arg](https://docs.bazel.build/versions/4.0.0/command-line-reference.html#flag--test_arg) flag. + E.g. `bazel test //src:rust_test --test_arg=foo::test::test_fn`. + """), + ), "_test_sharding_wrapper_unix": attr.label( default = Label("//rust/private:test_sharding_wrapper.sh"), allow_single_file = True, diff --git a/rust/private/test_sharding_wrapper.bat b/rust/private/test_sharding_wrapper.bat index 5c90681c81..e90a803f1c 100644 --- a/rust/private/test_sharding_wrapper.bat +++ b/rust/private/test_sharding_wrapper.bat @@ -14,7 +14,8 @@ @REM Wrapper script for rust_test that enables Bazel test sharding support. @REM This script intercepts test execution, enumerates tests using libtest's -@REM --list flag, partitions them by shard index, and runs only the relevant subset. +@REM --list flag, partitions them by stable test-name hash, and runs only the +@REM relevant subset. @ECHO OFF SETLOCAL EnableDelayedExpansion @@ -68,45 +69,118 @@ IF !FOUND_BINARY! EQU 0 IF DEFINED RUNFILES_DIR ( ) ) +@REM Try 4: manifest-based runfile lookup. This covers nested launchers that +@REM execute the sharding wrapper from another test's runfiles tree. +IF !FOUND_BINARY! EQU 0 ( + SET "MANIFEST=!RUNFILES_MANIFEST_FILE!" + IF NOT DEFINED MANIFEST IF EXIST "%~f0.runfiles_manifest" SET "MANIFEST=%~f0.runfiles_manifest" + IF NOT DEFINED MANIFEST IF EXIST "%~dpn0.runfiles_manifest" SET "MANIFEST=%~dpn0.runfiles_manifest" + IF NOT DEFINED MANIFEST IF EXIST "%~f0.exe.runfiles_manifest" SET "MANIFEST=%~f0.exe.runfiles_manifest" + + IF DEFINED MANIFEST IF EXIST "!MANIFEST!" ( + SET "TEST_BINARY_MANIFEST_PATH=!TEST_BINARY_RAW!" + SET "TEST_BINARY_MANIFEST_PATH=!TEST_BINARY_MANIFEST_PATH:\=/!" + IF DEFINED TEST_WORKSPACE SET "TEST_BINARY_MANIFEST_WORKSPACE_PATH=!TEST_WORKSPACE!/!TEST_BINARY_MANIFEST_PATH!" + FOR /F "usebackq tokens=1,* delims= " %%A IN ("!MANIFEST!") DO ( + IF "%%A"=="!TEST_BINARY_MANIFEST_PATH!" ( + SET "TEST_BINARY_PATH=%%B" + SET FOUND_BINARY=1 + GOTO :FOUND_TEST_BINARY + ) + IF DEFINED TEST_BINARY_MANIFEST_WORKSPACE_PATH IF "%%A"=="!TEST_BINARY_MANIFEST_WORKSPACE_PATH!" ( + SET "TEST_BINARY_PATH=%%B" + SET FOUND_BINARY=1 + GOTO :FOUND_TEST_BINARY + ) + ) + ) +) + +:FOUND_TEST_BINARY + IF !FOUND_BINARY! EQU 0 ( ECHO ERROR: Could not find test binary at any expected location EXIT /B 1 ) +@REM Native Bazel test sharding sets TEST_TOTAL_SHARDS/TEST_SHARD_INDEX. +@REM Explicit shard test targets can set RULES_RUST_TEST_TOTAL_SHARDS/ +@REM RULES_RUST_TEST_SHARD_INDEX instead because Bazel may reserve TEST_* +@REM variables for its own test runner env. +SET TOTAL_SHARDS=%RULES_RUST_TEST_TOTAL_SHARDS% +IF "%TOTAL_SHARDS%"=="" SET TOTAL_SHARDS=%TEST_TOTAL_SHARDS% +SET SHARD_INDEX=%RULES_RUST_TEST_SHARD_INDEX% +IF "%SHARD_INDEX%"=="" SET SHARD_INDEX=%TEST_SHARD_INDEX% + @REM If sharding is not enabled, run test binary directly -IF "%TEST_TOTAL_SHARDS%"=="" ( +IF "%TOTAL_SHARDS%"=="" ( + !TEST_BINARY_PATH! %* + EXIT /B !ERRORLEVEL! +) +IF "%TOTAL_SHARDS%"=="0" ( !TEST_BINARY_PATH! %* EXIT /B !ERRORLEVEL! ) +IF "%SHARD_INDEX%"=="" ( + ECHO ERROR: TEST_SHARD_INDEX or RULES_RUST_TEST_SHARD_INDEX must be set when sharding is enabled + EXIT /B 1 +) + @REM Touch status file to advertise sharding support to Bazel -IF NOT "%TEST_SHARD_STATUS_FILE%"=="" ( +IF NOT "%TEST_SHARD_STATUS_FILE%"=="" IF NOT "%TEST_TOTAL_SHARDS%"=="" IF NOT "%TEST_TOTAL_SHARDS%"=="0" ( TYPE NUL > "%TEST_SHARD_STATUS_FILE%" ) -@REM Create a temporary file for test list -SET TEMP_LIST=%TEMP%\rust_test_list_%RANDOM%.txt +@REM Create per-wrapper temporary files. Prefer Bazel's per-test temp directory; +@REM when falling back to the shared temp directory, avoid %RANDOM%-only file +@REM names that can collide across concurrently running Windows test shards. +SET "TEMP_ROOT=%TEST_TMPDIR%" +IF NOT DEFINED TEMP_ROOT SET "TEMP_ROOT=%TEMP%" +IF NOT DEFINED TEMP_ROOT SET "TEMP_ROOT=." +:CREATE_TEMP_DIR +SET "TEMP_DIR=!TEMP_ROOT!\rust_test_sharding_!RANDOM!_!RANDOM!_!RANDOM!" +MKDIR "!TEMP_DIR!" 2>NUL +IF ERRORLEVEL 1 GOTO :CREATE_TEMP_DIR +SET "TEMP_LIST=!TEMP_DIR!\list.txt" +SET "TEMP_SHARD_LIST=!TEMP_DIR!\shard.txt" @REM Enumerate all tests using libtest's --list flag !TEST_BINARY_PATH! --list --format terse 2>NUL > "!TEMP_LIST!" +IF ERRORLEVEL 1 ( + RMDIR /S /Q "!TEMP_DIR!" 2>NUL + EXIT /B 1 +) + +@REM Sort tests by ordinal name and filter this shard by stable FNV-1a hash so +@REM adding or removing one test does not move unrelated tests between shards. +@REM In the PowerShell fragment below, 2166136261 is the 32-bit FNV offset basis, +@REM 16777619 is the FNV prime, and 4294967295 is the UInt32 mask. Use decimal +@REM constants because Windows PowerShell can interpret 0xffffffff as -1. +powershell.exe -NoProfile -ExecutionPolicy Bypass -Command ^ + "$ErrorActionPreference = 'Stop';" ^ + "$tests = @(Get-Content -LiteralPath $env:TEMP_LIST | Where-Object { $_.EndsWith(': test') } | ForEach-Object { $_.Substring(0, $_.Length - 6) });" ^ + "[Array]::Sort($tests, [StringComparer]::Ordinal);" ^ + "$totalShards = [uint32]$env:TOTAL_SHARDS; $shardIndex = [uint32]$env:SHARD_INDEX;" ^ + "$fnvPrime = [uint64]16777619; $u32Mask = [uint64]4294967295;" ^ + "foreach ($test in $tests) { $hash = [uint32]2166136261; foreach ($byte in [Text.Encoding]::UTF8.GetBytes($test)) { $hash = [uint32](([uint64]($hash -bxor $byte) * $fnvPrime) -band $u32Mask) }; if (($hash %% $totalShards) -eq $shardIndex) { $test } }" ^ + > "!TEMP_SHARD_LIST!" +IF ERRORLEVEL 1 ( + RMDIR /S /Q "!TEMP_DIR!" 2>NUL + EXIT /B 1 +) -@REM Count tests and filter for this shard -SET INDEX=0 SET SHARD_TESTS= -FOR /F "tokens=1 delims=:" %%T IN ('TYPE "!TEMP_LIST!" ^| FINDSTR /E ": test"') DO ( - SET /A MOD=!INDEX! %% %TEST_TOTAL_SHARDS% - IF !MOD! EQU %TEST_SHARD_INDEX% ( - IF "!SHARD_TESTS!"=="" ( - SET SHARD_TESTS=%%T - ) ELSE ( - SET SHARD_TESTS=!SHARD_TESTS! %%T - ) +FOR /F "usebackq delims=" %%T IN ("!TEMP_SHARD_LIST!") DO ( + IF "!SHARD_TESTS!"=="" ( + SET SHARD_TESTS=%%T + ) ELSE ( + SET SHARD_TESTS=!SHARD_TESTS! %%T ) - SET /A INDEX=!INDEX! + 1 ) -DEL "!TEMP_LIST!" 2>NUL +RMDIR /S /Q "!TEMP_DIR!" 2>NUL @REM If no tests for this shard, exit successfully IF "!SHARD_TESTS!"=="" ( diff --git a/rust/private/test_sharding_wrapper.sh b/rust/private/test_sharding_wrapper.sh index e05970ba0a..b1f0fb55d7 100644 --- a/rust/private/test_sharding_wrapper.sh +++ b/rust/private/test_sharding_wrapper.sh @@ -15,40 +15,70 @@ # Wrapper script for rust_test that enables Bazel test sharding support. # This script intercepts test execution, enumerates tests using libtest's -# --list flag, partitions them by shard index, and runs only the relevant subset. +# --list flag, partitions them by stable test-name hash, and runs only the +# relevant subset. set -euo pipefail TEST_BINARY="{{TEST_BINARY}}" +# Native Bazel test sharding sets TEST_TOTAL_SHARDS/TEST_SHARD_INDEX. Explicit +# shard test targets can set RULES_RUST_TEST_TOTAL_SHARDS/RULES_RUST_TEST_SHARD_INDEX +# instead because Bazel may reserve TEST_* variables for its own test runner env. +TOTAL_SHARDS="${RULES_RUST_TEST_TOTAL_SHARDS:-${TEST_TOTAL_SHARDS:-}}" +SHARD_INDEX="${RULES_RUST_TEST_SHARD_INDEX:-${TEST_SHARD_INDEX:-}}" + +test_shard_index() { + local test_name="$1" + # FNV-1a 32-bit hash. The initial value is the FNV offset basis, and + # 16777619 is the FNV prime. This gives a stable, cheap string hash without + # depending on platform-specific tools being present in the test sandbox. + local hash=2166136261 + local byte + local char + local i + local LC_ALL=C + + for ((i = 0; i < ${#test_name}; i++)); do + char="${test_name:i:1}" + printf -v byte "%d" "'$char" + hash=$(( ((hash ^ byte) * 16777619) & 0xffffffff )) + done + + echo $(( hash % TOTAL_SHARDS )) +} # If sharding is not enabled, run test binary directly -if [[ -z "${TEST_TOTAL_SHARDS:-}" ]]; then +if [[ -z "${TOTAL_SHARDS}" || "${TOTAL_SHARDS}" == "0" ]]; then exec "./${TEST_BINARY}" "$@" fi +if [[ -z "${SHARD_INDEX}" ]]; then + echo "TEST_SHARD_INDEX or RULES_RUST_TEST_SHARD_INDEX must be set when sharding is enabled" >&2 + exit 1 +fi + # Touch status file to advertise sharding support to Bazel -if [[ -n "${TEST_SHARD_STATUS_FILE:-}" ]]; then +if [[ -n "${TEST_SHARD_STATUS_FILE:-}" && "${TEST_TOTAL_SHARDS:-0}" != "0" ]]; then touch "${TEST_SHARD_STATUS_FILE}" fi -# Enumerate all tests using libtest's --list flag +# Enumerate all tests using libtest's --list flag. Sort the list so execution +# order does not depend on libtest's output order. # Output format: "test_name: test" - we need to strip the ": test" suffix -test_list=$("./${TEST_BINARY}" --list --format terse 2>/dev/null | grep ': test$' | sed 's/: test$//' || true) +test_list=$("./${TEST_BINARY}" --list --format terse 2>/dev/null | grep ': test$' | sed 's/: test$//' | LC_ALL=C sort || true) # If no tests found, exit successfully if [[ -z "$test_list" ]]; then exit 0 fi -# Filter tests for this shard -# test_index % TEST_TOTAL_SHARDS == TEST_SHARD_INDEX +# Filter tests for this shard. Use a stable name hash instead of list position +# so adding or removing one test does not move unrelated tests between shards. shard_tests=() -index=0 while IFS= read -r test_name; do - if (( index % TEST_TOTAL_SHARDS == TEST_SHARD_INDEX )); then + if (( $(test_shard_index "$test_name") == SHARD_INDEX )); then shard_tests+=("$test_name") fi - ((index++)) || true done <<< "$test_list" # If no tests for this shard, exit successfully diff --git a/test/unit/test_sharding/fake_libtest_binary.sh b/test/unit/test_sharding/fake_libtest_binary.sh new file mode 100755 index 0000000000..fef029be5c --- /dev/null +++ b/test/unit/test_sharding/fake_libtest_binary.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash + +set -euo pipefail + +emit_base_tests() { + cat <<'EOF' +delta::test_delta: test +alpha::test_alpha: test +foxtrot::test_foxtrot: test +bravo::test_bravo: test +echo::test_echo: test +charlie::test_charlie: test +helper::bench: bench +EOF +} + +emit_reversed_base_tests() { + cat <<'EOF' +helper::bench: bench +charlie::test_charlie: test +echo::test_echo: test +bravo::test_bravo: test +foxtrot::test_foxtrot: test +alpha::test_alpha: test +delta::test_delta: test +EOF +} + +emit_tests_with_added_test() { + cat <<'EOF' +delta::test_delta: test +alpha::test_alpha: test +foxtrot::test_foxtrot: test +aardvark::test_added: test +bravo::test_bravo: test +echo::test_echo: test +charlie::test_charlie: test +helper::bench: bench +EOF +} + +if [[ "${1:-}" == "--list" ]]; then + case "${TEST_LIST_VARIANT:-base}:${TEST_LIST_ORDER:-normal}" in + base:normal) + emit_base_tests + ;; + base:reversed) + emit_reversed_base_tests + ;; + with_added:normal) + emit_tests_with_added_test + ;; + *) + echo "unknown test list variant: ${TEST_LIST_VARIANT:-base}:${TEST_LIST_ORDER:-normal}" >&2 + exit 1 + ;; + esac + exit 0 +fi + +: "${TEST_SHARD_OUTPUT:?}" + +for test_name in "$@"; do + if [[ "$test_name" != "--exact" ]]; then + printf '%s\n' "$test_name" >> "$TEST_SHARD_OUTPUT" + fi +done diff --git a/test/unit/test_sharding/test_sharding.bzl b/test/unit/test_sharding/test_sharding.bzl index aea76f170e..14005f7160 100644 --- a/test/unit/test_sharding/test_sharding.bzl +++ b/test/unit/test_sharding/test_sharding.bzl @@ -1,6 +1,7 @@ """Tests for rust_test sharding support.""" load("@bazel_skylib//lib:unittest.bzl", "analysistest", "asserts") +load("@rules_shell//shell:sh_test.bzl", "sh_test") load("//rust:defs.bzl", "rust_test") def _sharding_enabled_test(ctx): @@ -68,6 +69,23 @@ def _test_sharding_targets(): shard_count = 3, ) + sh_test( + name = "test_sharding_wrapper_hashes_sorted_names", + srcs = ["test_sharding_wrapper_hashes_sorted_names.sh"], + args = [ + "$(location //rust/private:test_sharding_wrapper.sh)", + "$(location :fake_libtest_binary.sh)", + ], + data = [ + ":fake_libtest_binary.sh", + "//rust/private:test_sharding_wrapper.sh", + ], + target_compatible_with = select({ + "@platforms//os:windows": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + ) + def test_sharding_test_suite(name): _test_sharding_targets() @@ -76,6 +94,7 @@ def test_sharding_test_suite(name): tests = [ ":sharding_enabled_test", ":sharding_disabled_test", + ":test_sharding_wrapper_hashes_sorted_names", ":sharded_integration_test", ], ) diff --git a/test/unit/test_sharding/test_sharding_wrapper_hashes_sorted_names.sh b/test/unit/test_sharding/test_sharding_wrapper_hashes_sorted_names.sh new file mode 100755 index 0000000000..626f33da44 --- /dev/null +++ b/test/unit/test_sharding/test_sharding_wrapper_hashes_sorted_names.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash + +set -euo pipefail + +wrapper_template=$1 +fake_binary_src=$2 + +workdir="${TEST_TMPDIR:-$(mktemp -d)}" +fake_binary="$workdir/fake_libtest_binary" +wrapper="$workdir/wrapper.sh" + +cp "$fake_binary_src" "$fake_binary" +chmod +x "$fake_binary" + +sed 's|{{TEST_BINARY}}|fake_libtest_binary|g' "$wrapper_template" > "$wrapper" +chmod +x "$wrapper" + +collect_mapping() { + local variant=$1 + local order=$2 + local output=$3 + local unsorted_output="${output}.unsorted" + local shard + + : > "$unsorted_output" + for shard in 0 1 2; do + local shard_output="$workdir/${variant}_${order}_${shard}.txt" + : > "$shard_output" + + ( + cd "$workdir" + TEST_LIST_VARIANT="$variant" \ + TEST_LIST_ORDER="$order" \ + TEST_SHARD_OUTPUT="$shard_output" \ + RULES_RUST_TEST_SHARD_INDEX="$shard" \ + RULES_RUST_TEST_TOTAL_SHARDS=3 \ + ./wrapper.sh + ) + + while IFS= read -r test_name; do + printf '%s %s\n' "$test_name" "$shard" >> "$unsorted_output" + done < "$shard_output" + done + + LC_ALL=C sort "$unsorted_output" > "$output" +} + +assert_same_mapping() { + local expected=$1 + local actual=$2 + local message=$3 + + if ! diff -u "$expected" "$actual"; then + echo "$message" >&2 + exit 1 + fi +} + +base_normal="$workdir/base_normal.txt" +base_reversed="$workdir/base_reversed.txt" +with_added="$workdir/with_added.txt" +with_added_existing_tests="$workdir/with_added_existing_tests.txt" + +collect_mapping base normal "$base_normal" +collect_mapping base reversed "$base_reversed" +collect_mapping with_added normal "$with_added" + +assert_same_mapping "$base_normal" "$base_reversed" \ + "test shard assignment changed when libtest list order changed" + +sed '/^aardvark::test_added /d' "$with_added" > "$with_added_existing_tests" +assert_same_mapping "$base_normal" "$with_added_existing_tests" \ + "existing test shard assignment changed after adding a new test"