From c8db1518cfd86e5f4e953c768a99fd7d4410411e Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 27 Sep 2023 14:20:56 -0700
Subject: [PATCH] build: add roachtest nightlies coverage scripts

This script will be used in a coverage build that runs the nightlies
on GCE and collects coverage data.

Epic: none
Release note: None
---
 .../coverage/roachtest_nightly_gce.sh         | 12 ++++++
 .../coverage/roachtest_nightly_gce_impl.sh    | 39 +++++++++++++++++++
 .../nightlies/roachtest_compile_bits.sh       | 30 +++++++++++---
 3 files changed, 75 insertions(+), 6 deletions(-)
 create mode 100755 build/teamcity/cockroach/coverage/roachtest_nightly_gce.sh
 create mode 100755 build/teamcity/cockroach/coverage/roachtest_nightly_gce_impl.sh

diff --git a/build/teamcity/cockroach/coverage/roachtest_nightly_gce.sh b/build/teamcity/cockroach/coverage/roachtest_nightly_gce.sh
new file mode 100755
index 000000000000..253daf668dc0
--- /dev/null
+++ b/build/teamcity/cockroach/coverage/roachtest_nightly_gce.sh
@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+
+source "$dir/teamcity-support.sh"  # For $root
+source "$dir/teamcity-bazel-support.sh"  # For run_bazel
+
+CLOUD=gce
+BAZEL_SUPPORT_EXTRA_DOCKER_ARGS="-e LITERAL_ARTIFACTS_DIR=$root/artifacts -e BUILD_VCS_NUMBER -e CLOUD -e COCKROACH_DEV_LICENSE -e TESTS -e COUNT -e GITHUB_API_TOKEN -e GITHUB_ORG -e GITHUB_REPO -e GOOGLE_EPHEMERAL_CREDENTIALS -e GOOGLE_KMS_KEY_A -e GOOGLE_KMS_KEY_B -e GOOGLE_CREDENTIALS_ASSUME_ROLE -e GOOGLE_SERVICE_ACCOUNT -e SLACK_TOKEN -e TC_BUILDTYPE_ID -e TC_BUILD_BRANCH -e TC_BUILD_ID -e TC_SERVER_URL -e SELECT_PROBABILITY" \
+			       run_bazel build/teamcity/cockroach/coverage/roachtest_nightly_gce_impl.sh
diff --git a/build/teamcity/cockroach/coverage/roachtest_nightly_gce_impl.sh b/build/teamcity/cockroach/coverage/roachtest_nightly_gce_impl.sh
new file mode 100755
index 000000000000..5f4f1c5088db
--- /dev/null
+++ b/build/teamcity/cockroach/coverage/roachtest_nightly_gce_impl.sh
@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+
+source "$dir/teamcity-support.sh"
+
+if [[ ! -f ~/.ssh/id_rsa.pub ]]; then
+  ssh-keygen -q -C "roachtest-nightly-bazel $(date)" -N "" -f ~/.ssh/id_rsa
+fi
+
+# We will only spin up AMD64 clusters. Note that the TeamCity runner (this host)
+# doesn't need to be AMD64; the roachtest binary is always built for the host
+# architecture.
+$root/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh --with-code-coverage amd64
+
+echo "$GOOGLE_EPHEMERAL_CREDENTIALS" > creds.json
+gcloud auth activate-service-account --key-file=creds.json
+export ROACHPROD_USER=teamcity
+
+# See build/teamcity/util/roachtest_util.sh.
+PARALLELISM=16
+CPUQUOTA=1024
+FILTER="tag:aws tag:default"
+
+build/teamcity-roachtest-invoke.sh \
+  --metamorphic-encryption-probability=0.5 \
+  --select-probability="${SELECT_PROBABILITY:-1.0}" \
+  --cloud="${CLOUD}" \
+  --count="${COUNT-1}" \
+  --parallelism="${PARALLELISM}" \
+  --cpu-quota="${CPUQUOTA}" \
+  --cluster-id="${TC_BUILD_ID}" \
+  --artifacts=/artifacts \
+  --artifacts-literal="${LITERAL_ARTIFACTS_DIR:-}" \
+  --slack-token="${SLACK_TOKEN:-}" \
+  --go-cover \
+  ${TESTS:-} ${FILTER}
diff --git a/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh b/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
index 6fc4938cf185..c6ff224a73ab 100755
--- a/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
+++ b/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
@@ -5,7 +5,7 @@ set -euo pipefail
 if [ "$#" -eq 0 ]; then
   echo "Builds all bits needed for roachtests and stages them in bin/ and lib/."
   echo ""
-  echo "Usage: $0 arch [arch...]"
+  echo "Usage: $0 [--with-code-coverage] arch [arch...]"
   echo "  where arch is one of: amd64, arm64, amd64-fips"
   exit 1
 fi
@@ -29,12 +29,29 @@ function arch_to_config() {
   esac
 }
 
+arches=()
+crdb_extra_flags=""
+
+for arg in "$@"; do
+  case "$arg" in
+    --with-code-coverage)
+      crdb_extra_flags="--collect_code_coverage --bazel_code_coverage"
+      ;;
+    *)
+      # Fail now if the argument is not a valid arch.
+      arch_to_config $arg >/dev/null || exit 1
+      arches+=($arg)
+      ;;
+  esac
+done
+
 # Determine host cpu architecture, which we'll need for libgeos, below.
 if [[ "$(uname -m)" =~ (arm64|aarch64)$ ]]; then
   host_arch=arm64
 else
   host_arch=amd64
 fi
+
 echo "Host architecture: $host_arch"
 
 # Prepare the bin/ and lib/ directories.
@@ -43,17 +60,18 @@ chmod o+rwx bin
 mkdir -p lib
 chmod o+rwx lib
 
-for arch in "$@"; do
+for arch in "${arches[@]}"; do
   config=$(arch_to_config $arch)
   echo "Building $config, os=$os, arch=$arch..."
   # Build cockroach, workload and geos libs.
   bazel build --config $config --config ci -c opt --config force_build_cdeps \
-        //pkg/cmd/cockroach //pkg/cmd/workload \
-        //c-deps:libgeos
+        //pkg/cmd/cockroach $crdb_extra_flags
+  bazel build --config $config --config ci -c opt --config force_build_cdeps \
+        //pkg/cmd/workload //c-deps:libgeos
   BAZEL_BIN=$(bazel info bazel-bin --config $config --config ci -c opt)
 
-  # Build cockroach-short with assertions enabled.
-  bazel build --config $config --config ci -c opt //pkg/cmd/cockroach-short --crdb_test
+  bazel build --config $config --config ci -c opt //pkg/cmd/cockroach-short \
+        --crdb_test $crdb_extra_flags
   # Copy the binaries.
   cp $BAZEL_BIN/pkg/cmd/cockroach/cockroach_/cockroach bin/cockroach.$os-$arch
   cp $BAZEL_BIN/pkg/cmd/workload/workload_/workload    bin/workload.$os-$arch