bazel: run wrapped Rust unit test shards

defs.bzl

@@ -255,10 +255,9 @@ def codex_rust_crate(
         unit_test_name = name + "-unit-tests"
         unit_test_binary = name + "-unit-tests-bin"
         unit_test_shard_count = _test_shard_count(test_shard_counts, unit_test_name)
-        unit_test_binary_kwargs = {}
-        if unit_test_shard_count:
-            unit_test_binary_kwargs["experimental_enable_sharding"] = True
-
+        # Shard at the workspace_root_test layer. rules_rust's sharding wrapper
+        # expects to run from its own runfiles cwd, while workspace_root_test
+        # deliberately changes cwd so Insta sees Cargo-like snapshot paths.
         rust_test(
             name = unit_test_binary,
             crate = name,
@@ -277,7 +276,6 @@ def codex_rust_crate(
             rustc_env = rustc_env,
             data = test_data_extra,
             tags = test_tags + ["manual"],
-            **unit_test_binary_kwargs
         )
 
         unit_test_kwargs = {}

workspace_root_test_launcher.bat.tpl

@@ -12,9 +12,76 @@ if errorlevel 1 exit /b 1
 
 set "INSTA_WORKSPACE_ROOT=%workspace_root%"
 cd /d "%workspace_root%" || exit /b 1
+
+set "TOTAL_SHARDS=%RULES_RUST_TEST_TOTAL_SHARDS%"
+if not defined TOTAL_SHARDS set "TOTAL_SHARDS=%TEST_TOTAL_SHARDS%"
+if defined TOTAL_SHARDS if not "%TOTAL_SHARDS%"=="0" (
+  call :run_sharded_libtest %*
+  exit /b !ERRORLEVEL!
+)
+
 "%test_bin%" %*
 exit /b %ERRORLEVEL%
 
+:run_sharded_libtest
+if defined TEST_SHARD_STATUS_FILE if defined TEST_TOTAL_SHARDS if not "%TEST_TOTAL_SHARDS%"=="0" (
+  type nul > "%TEST_SHARD_STATUS_FILE%"
+)
+
+if not "%~1"=="" (
+  "%test_bin%" %*
+  exit /b !ERRORLEVEL!
+)
+
+set "SHARD_INDEX=%RULES_RUST_TEST_SHARD_INDEX%"
+if not defined SHARD_INDEX set "SHARD_INDEX=%TEST_SHARD_INDEX%"
+if not defined SHARD_INDEX (
+  >&2 echo TEST_SHARD_INDEX or RULES_RUST_TEST_SHARD_INDEX must be set when sharding is enabled
+  exit /b 1
+)
+
+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%\workspace_root_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"
+
+"%test_bin%" --list --format terse > "!TEMP_LIST!"
+if errorlevel 1 (
+  rmdir /s /q "!TEMP_DIR!" 2>nul
+  exit /b 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
+)
+
+powershell.exe -NoProfile -ExecutionPolicy Bypass -Command ^
+  "$ErrorActionPreference = 'Stop';" ^
+  "$testBin = $env:test_bin;" ^
+  "$tests = @(Get-Content -LiteralPath $env:TEMP_SHARD_LIST);" ^
+  "$failed = $false; $limit = 7000; $batch = @(); $batchChars = $testBin.Length + 8;" ^
+  "function Invoke-TestBatch { if ($script:batch.Count -eq 0) { return }; & $script:testBin @script:batch '--exact'; if ($LASTEXITCODE -ne 0) { $script:failed = $true }; $script:batch = @(); $script:batchChars = $script:testBin.Length + 8 }" ^
+  "foreach ($test in $tests) { $argChars = $test.Length + 3; if (($batch.Count -gt 0) -and ($batchChars + $argChars -gt $limit)) { Invoke-TestBatch }; $batch += $test; $batchChars += $argChars }" ^
+  "Invoke-TestBatch; if ($failed) { exit 1 }"
+set "TEST_EXIT=%ERRORLEVEL%"
+
+rmdir /s /q "!TEMP_DIR!" 2>nul
+exit /b !TEST_EXIT!
+
 :resolve_runfile
 setlocal EnableExtensions EnableDelayedExpansion
 set "logical_path=%~2"

workspace_root_test_launcher.sh.tpl

@@ -48,6 +48,72 @@ workspace_root_marker="$(resolve_runfile "__WORKSPACE_ROOT_MARKER__")"
 workspace_root="$(dirname "$(dirname "$(dirname "${workspace_root_marker}")")")"
 test_bin="$(resolve_runfile "__TEST_BIN__")"
 
+test_shard_index() {
+  local test_name="$1"
+  # FNV-1a 32-bit hash. Keep this stable so adding one test does not reshuffle
+  # unrelated tests between shards.
+  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 ))
+}
+
+run_sharded_libtest() {
+  if [[ -n "${TEST_SHARD_STATUS_FILE:-}" && "${TEST_TOTAL_SHARDS:-0}" != "0" ]]; then
+    touch "${TEST_SHARD_STATUS_FILE}"
+  fi
+
+  # Extra libtest args are usually ad-hoc local filters. Preserve those exactly
+  # rather than combining them with generated exact filters.
+  if [[ $# -gt 0 ]]; then
+    exec "${test_bin}" "$@"
+  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
+
+  local list_output
+  local test_list
+  list_output="$("${test_bin}" --list --format terse)"
+  test_list="$(printf '%s\n' "${list_output}" | grep ': test$' | sed 's/: test$//' | LC_ALL=C sort || true)"
+
+  if [[ -z "${test_list}" ]]; then
+    exit 0
+  fi
+
+  local shard_tests=()
+  local test_name
+  while IFS= read -r test_name; do
+    if (( $(test_shard_index "${test_name}") == SHARD_INDEX )); then
+      shard_tests+=("${test_name}")
+    fi
+  done <<< "${test_list}"
+
+  if [[ ${#shard_tests[@]} -eq 0 ]]; then
+    exit 0
+  fi
+
+  exec "${test_bin}" "${shard_tests[@]}" --exact
+}
+
 export INSTA_WORKSPACE_ROOT="${workspace_root}"
 cd "${workspace_root}"
+
+TOTAL_SHARDS="${RULES_RUST_TEST_TOTAL_SHARDS:-${TEST_TOTAL_SHARDS:-}}"
+SHARD_INDEX="${RULES_RUST_TEST_SHARD_INDEX:-${TEST_SHARD_INDEX:-}}"
+if [[ -n "${TOTAL_SHARDS}" && "${TOTAL_SHARDS}" != "0" ]]; then
+  run_sharded_libtest "$@"
+fi
+
 exec "${test_bin}" "$@"