From 7ac059f72bbbbe3cbfdf2485a7bc5b9eb61abda8 Mon Sep 17 00:00:00 2001
From: healthy-pod <ahmad@cockroachlabs.com>
Date: Tue, 6 Sep 2022 18:48:10 -0700
Subject: [PATCH] ci: pass custom timeout to testrace in CI

In #86363, we added a timeout to tests at the test binary level.
Tests running with `--config=race` however use a custom
timeout, different from the original default values set by bazel
based on the test size.

This patch propagates those custom values to testrace in CI.

Release justification: Non-production code changes
Release note: None
---
 .bazelrc                                      |  2 ++
 .../cockroach/ci/tests/testrace_impl.sh       | 30 ++++++++++++-------
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index 57ec9306980b..7731062479a6 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -34,6 +34,8 @@ info --ui_event_filters=-WARNING
 
 build:race --@io_bazel_rules_go//go/config:race "--test_env=GORACE=halt_on_error=1 log_path=stdout" --test_sharding_strategy=disabled
 test:test --test_env=TZ=
+# Note: these timeout values are used indirectly in `build/teamcity/cockroach/ci/tests/testrace_impl.sh`.
+# If those values are updated, the script should be updated accordingly.
 test:race --test_timeout=1200,6000,18000,72000
 
 # CI should always run with `--config=ci` or `--config=cinolint`.
diff --git a/build/teamcity/cockroach/ci/tests/testrace_impl.sh b/build/teamcity/cockroach/ci/tests/testrace_impl.sh
index 93b2475a7038..d34e54871e59 100755
--- a/build/teamcity/cockroach/ci/tests/testrace_impl.sh
+++ b/build/teamcity/cockroach/ci/tests/testrace_impl.sh
@@ -6,6 +6,7 @@ set -xeuo pipefail
 # packages are expected to be formatted as go-style, e.g. ./pkg/cmd/bazci.
 
 bazel build //pkg/cmd/bazci --config=ci
+size_to_timeout=("small:1200" "medium:6000" "large:18000" "enormous:72000")
 for pkg in "$@"
 do
     # Query to list all affected tests.
@@ -14,19 +15,26 @@ do
     then
         pkg="$pkg:all"
     fi
-    tests=$(bazel query "kind(go_test, $pkg)" --output=label)
 
-    # Run affected tests.
-    for test in $tests
+    for kv in "${size_to_timeout[@]}";
     do
-        if [[ ! -z $(bazel query "attr(tags, \"broken_in_bazel\", $test)") ]]
-        then
-            echo "Skipping test $test as it is broken in bazel"
-            continue
-        fi
-        $(bazel info bazel-bin --config=ci)/pkg/cmd/bazci/bazci_/bazci -- test --config=ci --config=race "$test" \
-                               --test_env=COCKROACH_LOGIC_TESTS_SKIP=true \
-                               --test_env=GOMAXPROCS=8
+        size="${kv%%:*}"
+        timeout="${kv#*:}"
+        go_timeout=$(($timeout - 5))
+        tests=$(bazel query "attr(size, $size, kind("go_test", tests($pkg)))" --output=label)
+        # Run affected tests.
+        for test in $tests
+        do
+            if [[ ! -z $(bazel query "attr(tags, \"broken_in_bazel\", $test)") ]]
+            then
+                echo "Skipping test $test as it is broken in bazel"
+                continue
+            fi
+            $(bazel info bazel-bin --config=ci)/pkg/cmd/bazci/bazci_/bazci -- test --config=ci --config=race "$test" \
+                                --test_env=COCKROACH_LOGIC_TESTS_SKIP=true \
+                                --test_env=GOMAXPROCS=8 \
+                                --test_arg=-test.timeout="${go_timeout}s"
+        done
     done
 done