From cd537e15b256678279beab7d4afbc93ac1ad75ad Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 4 Oct 2023 05:25:40 -0700
Subject: [PATCH] build: improve roachtest build scripts

We improve the roachtest scripts to support building finer-grained
components: the `roachtest_compile_component.sh` script builds a
single component (cockroach, cockroach-ea, workload, libgeos,
roachtest). For each component we define the bazel arguments and build
products, minimizing duplication of other code. There should be no
functional changes.

Epic: none
Release note: None
---
 .../nightlies/roachtest_arch_util.sh          |  27 +++++
 .../nightlies/roachtest_compile_bits.sh       |  95 ++++------------
 .../nightlies/roachtest_compile_component.sh  | 105 ++++++++++++++++++
 3 files changed, 154 insertions(+), 73 deletions(-)
 create mode 100644 build/teamcity/cockroach/nightlies/roachtest_arch_util.sh
 create mode 100755 build/teamcity/cockroach/nightlies/roachtest_compile_component.sh

diff --git a/build/teamcity/cockroach/nightlies/roachtest_arch_util.sh b/build/teamcity/cockroach/nightlies/roachtest_arch_util.sh
new file mode 100644
index 000000000000..c27de99d4854
--- /dev/null
+++ b/build/teamcity/cockroach/nightlies/roachtest_arch_util.sh
@@ -0,0 +1,27 @@
+# arch_to_config returns the bazel config for the given arch.
+function arch_to_config() {
+  case "$1" in
+    amd64)
+      echo "crosslinux"
+      ;;
+    arm64)
+      echo "crosslinuxarm"
+      ;;
+    amd64-fips)
+      echo "crosslinuxfips"
+      ;;
+    *)
+      echo "Error: invalid arch '$1'" >&2
+      exit 1
+      ;;
+  esac
+}
+
+# get_host_arch determines the current host's CPU architecture.
+function get_host_arch() {
+  if [[ "$(uname -m)" =~ (arm64|aarch64)$ ]]; then
+    echo "arm64"
+  else
+    echo "amd64"
+  fi
+}
diff --git a/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh b/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
index c6ff224a73ab..e36285d90fd0 100755
--- a/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
+++ b/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
@@ -2,6 +2,8 @@
 
 set -euo pipefail
 
+source "$(dirname $0)/roachtest_arch_util.sh"
+
 if [ "$#" -eq 0 ]; then
   echo "Builds all bits needed for roachtests and stages them in bin/ and lib/."
   echo ""
@@ -11,100 +13,47 @@ if [ "$#" -eq 0 ]; then
 fi
 
 os=linux
-function arch_to_config() {
-  case "$1" in
-    amd64)
-      echo "crosslinux"
-      ;;
-    arm64)
-      echo "crosslinuxarm"
-      ;;
-    amd64-fips)
-      echo "crosslinuxfips"
-      ;;
-    *)
-      echo "Error: invalid arch '$1'" >&2
-      exit 1
-      ;;
-  esac
-}
 
-arches=()
-crdb_extra_flags=""
+components=()
+extra_flags=""
 
 for arg in "$@"; do
   case "$arg" in
     --with-code-coverage)
-      crdb_extra_flags="--collect_code_coverage --bazel_code_coverage"
+      extra_flags="$arg"
       ;;
     *)
       # Fail now if the argument is not a valid arch.
       arch_to_config $arg >/dev/null || exit 1
-      arches+=($arg)
+      components+=($os/$arg/cockroach)
+      components+=($os/$arg/cockroach-ea)
+      components+=($os/$arg/workload)
+      components+=($os/$arg/libgeos)
       ;;
   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
-
+# We need to build roachtest and geos libraries (necessary for local tests) for
+# the host architecture.
+host_arch=$(get_host_arch)
 echo "Host architecture: $host_arch"
+components+=($os/$host_arch/roachtest)
+components+=($os/$host_arch/libgeos)
 
 # Prepare the bin/ and lib/ directories.
-mkdir -p bin
-chmod o+rwx bin
-mkdir -p lib
-chmod o+rwx lib
-
-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 $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)
-
-  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
-  # N.B. "-ea" suffix stands for enabled-assertions.
-  cp $BAZEL_BIN/pkg/cmd/cockroach-short/cockroach-short_/cockroach-short bin/cockroach-ea.$os-$arch
-  # Make it writable to simplify cleanup and copying (e.g., scp retry).
-  chmod a+w bin/cockroach.$os-$arch bin/workload.$os-$arch bin/cockroach-ea.$os-$arch
+mkdir -p bin lib
+chmod o+rwx bin lib
 
-  # Copy geos libs.
-  cp $BAZEL_BIN/c-deps/libgeos_foreign/lib/libgeos.so   lib/libgeos.$os-$arch.so
-  cp $BAZEL_BIN/c-deps/libgeos_foreign/lib/libgeos_c.so lib/libgeos_c.$os-$arch.so
-  # Make it writable to simplify cleanup and copying (e.g., scp retry).
-  chmod a+w lib/libgeos.$os-$arch.so lib/libgeos_c.$os-$arch.so
+# Sort and dedup components (libgeos can show up twice).
+for comp in $(printf "%s\n" "${components[@]}" | sort -u); do
+  "$(dirname $0)"/roachtest_compile_component.sh $extra_flags $comp
 done
 
-# Build roachtest itself, for the host architecture.
-# We also build geos libs; roachtest runner may require them locally.
-config=$(arch_to_config $host_arch)
-echo "Building roachtest and libgeos ($config, os=$os, arch=$host_arch)..."
-
-bazel build --config $config --config ci  -c opt //pkg/cmd/roachtest
-bazel build --config $config --config ci -c opt --config force_build_cdeps \
-      //c-deps:libgeos
-
-BAZEL_BIN=$(bazel info bazel-bin --config $config --config ci -c opt)
-cp $BAZEL_BIN/pkg/cmd/roachtest/roachtest_/roachtest bin/roachtest
+cp -p bin/roachtest.$os-$host_arch bin/roachtest
 # N.B. geos does not support the architecture suffix (see getLibraryExt() in
 # geos.go).
-cp $BAZEL_BIN/c-deps/libgeos_foreign/lib/libgeos.so   lib/libgeos.so
-cp $BAZEL_BIN/c-deps/libgeos_foreign/lib/libgeos_c.so lib/libgeos_c.so
-
-# Make files writable to simplify cleanup and copying (e.g., scp retry).
-chmod a+w bin/roachtest lib/libgeos.so lib/libgeos_c.so
+cp -p lib/libgeos.$os-$host_arch.so lib/libgeos.so
+cp -p lib/libgeos_c.$os-$host_arch.so lib/libgeos_c.so
 
 ls -l bin
 ls -l lib
diff --git a/build/teamcity/cockroach/nightlies/roachtest_compile_component.sh b/build/teamcity/cockroach/nightlies/roachtest_compile_component.sh
new file mode 100755
index 000000000000..fcae0cab5576
--- /dev/null
+++ b/build/teamcity/cockroach/nightlies/roachtest_compile_component.sh
@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+source "$(dirname $0)/roachtest_arch_util.sh"
+
+if [ "$#" -eq 0 ]; then
+  echo "Builds components necessary for roachtests and stages them in bin/ and/or lib/."
+  echo ""
+  echo "Usage: $(basename $0) [--with-coverage] <os/arch/component>"
+  echo "  where os is one of: linux"
+  echo "        arch is one of: amd64, arm64, amd64-fips"
+  echo "        component is one of: cockroach, cockroach-ea, workload, libgeos, roachtest"
+  echo "  --with-coverage enables go code coverage instrumentation (only applies to cockroach binaries)"
+  exit 1
+fi
+
+crdb_extra_flags=""
+target=""
+
+for arg in "$@"; do
+  case "$arg" in
+    --with-code-coverage)
+      crdb_extra_flags="--collect_code_coverage --bazel_code_coverage"
+      ;;
+    *)
+      if [ -n "$target" ]; then
+        echo "Error: too many arguments"
+        exit 1
+      fi
+      target=$arg
+      ;;
+  esac
+done
+
+# Split $target into $os/$arch/$component.
+os=${target%%/*}
+arch_and_comp=${target#*/}
+arch=${arch_and_comp%%/*}
+component=${arch_and_comp#*/}
+
+if [ -z $os -o -z $arch -o -z $component ]; then
+  echo "Invalid target: $target; syntax is os/arch/component"
+  exit 1
+fi
+
+if [ "$os" != linux ]; then
+  echo "Invalid os; only linux supported"
+  exit 1
+fi
+
+config=$(arch_to_config $arch)
+
+# Array of arguments to be passed to bazel for the component.
+bazel_args=()
+# Array of build artifacts. Each item has format "src:dest"; src is relative to
+# the bazel-bin directory, dst is relative to cwd.
+artifacts=()
+
+case "$component" in
+  cockroach)
+    # Cockroach binary.
+    bazel_args=(--config force_build_cdeps //pkg/cmd/cockroach $crdb_extra_flags)
+    artifacts=("pkg/cmd/cockroach/cockroach_/cockroach:bin/cockroach.$os-$arch")
+    ;;
+  cockroach-ea)
+    # Cockroach-short with enabled assertions (EA).
+    bazel_args=(--config force_build_cdeps //pkg/cmd/cockroach-short --crdb_test $crdb_extra_flags)
+    artifacts=("pkg/cmd/cockroach-short/cockroach-short_/cockroach-short:bin/cockroach-ea.$os-$arch")
+    ;;
+  workload)
+    # Workload binary.
+    bazel_args=(--config force_build_cdeps //pkg/cmd/workload)
+    artifacts=("pkg/cmd/workload/workload_/workload:bin/workload.$os-$arch")
+    ;;
+  libgeos)
+    # Geos libs.
+    bazel_args=(--config force_build_cdeps //c-deps:libgeos)
+    artifacts=(
+      "c-deps/libgeos_foreign/lib/libgeos.so:lib/libgeos.$os-$arch.so"
+      "c-deps/libgeos_foreign/lib/libgeos_c.so:lib/libgeos_c.$os-$arch.so"
+    )
+    ;;
+  roachtest)
+    # Roachtest binary.
+    bazel_args=(//pkg/cmd/roachtest)
+    artifacts=("pkg/cmd/roachtest/roachtest_/roachtest:bin/roachtest.$os-$arch")
+    ;;
+  *)
+    echo "Unknown component '$component'"
+    exit 1
+    ;;
+esac
+
+echo "Building $os/$arch/$component..."
+
+bazel build --config $config --config ci -c opt "${bazel_args[@]}"
+BAZEL_BIN=$(bazel info bazel-bin --config $config --config ci -c opt)
+for artifact in "${artifacts[@]}"; do
+  src=${artifact%%:*}
+  dst=${artifact#*:}
+  cp "$BAZEL_BIN/$src" "$dst"
+  # Make files writable to simplify cleanup and copying (e.g., scp retry).
+  chmod a+w "$dst"
+done