Merge #87760

87760: generate-bazel-extra: optimize timeouts generation r=rickystewart a=healthy-pod

This code change reduces the time it takes to generate timeouts for go tests. Previously, it took an average of 17 seconds to generate timeouts. Now it takes an average of 5 seconds to do so.

An average of 2 seconds was reduced by avoiding `pkg/ui` subdirectories and an average of 10 seconds was reduced by running `bazel query` once instead of running it 4 times.

Release note: None

Co-authored-by: healthy-pod <[email protected]>

pkg/BUILD.bazel

@@ -116,6 +116,7 @@ ALL_TESTS = [
     "//pkg/cmd/dev:dev_test",
     "//pkg/cmd/docgen/extract:extract_test",
     "//pkg/cmd/docs-issue-generation:docs-issue-generation_test",
+    "//pkg/cmd/generate-bazel-extra:generate-bazel-extra_test",
     "//pkg/cmd/github-post:github-post_test",
     "//pkg/cmd/github-pull-request-make:github-pull-request-make_test",
     "//pkg/cmd/internal/issues:issues_test",
@@ -520,6 +521,7 @@ ALL_TESTS = [
     "//pkg/testutils:testutils_test",
     "//pkg/ts/testmodel:testmodel_test",
     "//pkg/ts:ts_test",
+    "//pkg/ui:ui_test",
     "//pkg/upgrade/upgradecluster:upgradecluster_test",
     "//pkg/upgrade/upgrademanager:upgrademanager_test",
     "//pkg/upgrade/upgrades:upgrades_test",
@@ -884,6 +886,7 @@ GO_TARGETS = [
     "//pkg/cmd/fuzz:fuzz_lib",
     "//pkg/cmd/generate-bazel-extra:generate-bazel-extra",
     "//pkg/cmd/generate-bazel-extra:generate-bazel-extra_lib",
+    "//pkg/cmd/generate-bazel-extra:generate-bazel-extra_test",
     "//pkg/cmd/generate-binary:generate-binary",
     "//pkg/cmd/generate-binary:generate-binary_lib",
     "//pkg/cmd/generate-distdir:generate-distdir",
@@ -1853,6 +1856,8 @@ GO_TARGETS = [
     "//pkg/ts/tspb:tspb",
     "//pkg/ts:ts",
     "//pkg/ts:ts_test",
+    "//pkg/ui:ui",
+    "//pkg/ui:ui_test",
     "//pkg/upgrade/nodelivenesstest:nodelivenesstest",
     "//pkg/upgrade/upgradecluster:upgradecluster",
     "//pkg/upgrade/upgradecluster:upgradecluster_test",
@@ -2866,6 +2871,7 @@ GET_X_DATA_TARGETS = [
     "//pkg/ts/catalog:get_x_data",
     "//pkg/ts/testmodel:get_x_data",
     "//pkg/ts/tspb:get_x_data",
+    "//pkg/ui:get_x_data",
     "//pkg/upgrade:get_x_data",
     "//pkg/upgrade/nodelivenesstest:get_x_data",
     "//pkg/upgrade/upgradecluster:get_x_data",

pkg/cmd/generate-bazel-extra/BUILD.bazel

@@ -1,5 +1,5 @@
 load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
-load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library", "go_test")
 
 go_library(
     name = "generate-bazel-extra_lib",
@@ -15,4 +15,17 @@ go_binary(
     visibility = ["//visibility:public"],
 )
 
+go_test(
+    name = "generate-bazel-extra_test",
+    size = "small",
+    srcs = ["main_test.go"],
+    args = ["-test.timeout=55s"],
+    data = glob(["testdata/**"]),
+    embed = [":generate-bazel-extra_lib"],
+    deps = [
+        "//pkg/testutils",
+        "@com_github_stretchr_testify//require",
+    ],
+)
+
 get_x_data(name = "get_x_data")

pkg/cmd/generate-bazel-extra/main.go

@@ -18,6 +18,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -48,15 +49,49 @@ func runBuildozer(args []string) {
 	}
 }
 
-func getTestTargets(testTargetSize string) ([]string, error) {
-	if _, ok := testSizeToDefaultTimeout[testTargetSize]; !ok {
-		return nil, errors.New("testTargetSize should be one of {small,medium,large,enormous}")
+// parseQueryXML is used because Go doesn't support parsing XML "1.1".
+// It returns a map where the key is a test size {small,medium,large,enormous} and
+// the value is a list of test targets having that size.
+func parseQueryXML(data []byte) (map[string][]string, error) {
+	targetNameRegex, err := regexp.Compile(`<rule.*class="go_test".*name="(.*)".*>`)
+	if err != nil {
+		return nil, err
+	}
+
+	targetSizeRegex, err := regexp.Compile(`<string.*name="size".*value="(.*)".*/>`)
+	if err != nil {
+		return nil, err
+	}
+
+	targetToSize := make(map[string]string)
+	var currentTargetName string
+	for _, line := range strings.Split(string(data), "\n") {
+		// Check if the line contains a target name.
+		line = strings.TrimSpace(line)
+		if submatch := targetNameRegex.FindStringSubmatch(line); submatch != nil {
+			currentTargetName = submatch[1]
+			// Default size is medium so if not found then it will be medium.
+			targetToSize[currentTargetName] = "medium"
+			continue
+		}
+		// Check if the line contains a target size.
+		if submatch := targetSizeRegex.FindStringSubmatch(line); submatch != nil {
+			targetToSize[currentTargetName] = submatch[1]
+		}
+	}
+	sizeToTargets := make(map[string][]string)
+	for target, size := range targetToSize {
+		sizeToTargets[size] = append(sizeToTargets[size], target)
 	}
+	return sizeToTargets, nil
+}
+
+func getTestTargets() (map[string][]string, error) {
 	cmd := exec.Command(
 		"bazel",
 		"query",
-		fmt.Sprintf(`attr(size, %s, kind("go_test", tests(//pkg/...)))`, testTargetSize),
-		"--output=label",
+		fmt.Sprintf(`kind("go_test", %s)`, getPackagesToQuery()),
+		"--output=xml",
 	)
 	buf, err := cmd.Output()
 	if err != nil {
@@ -69,28 +104,33 @@ func getTestTargets(testTargetSize string) ([]string, error) {
 		}
 		os.Exit(1)
 	}
-	return strings.Split(strings.TrimSpace(string(buf[:])), "\n"), nil
+	return parseQueryXML(buf)
 }
 
-func generateTestSuites() {
+func getPackagesToQuery() string {
 	// First list all test and binary targets.
 	infos, err := os.ReadDir("pkg")
 	if err != nil {
 		panic(err)
 	}
 	var packagesToQuery []string
 	for _, info := range infos {
-		// We don't want to query into pkg/ui because it never contains any Go tests and
-		// because doing so causes a pull from `npm`.
+		// We don't want to query into pkg/ui because it only contains a
+		// single Go test target at its root which will be included below.
+		// Querying into its subdirectories is unneeded and causes a pull from `npm`.
 		if !info.IsDir() || info.Name() == "ui" {
 			continue
 		}
 		packagesToQuery = append(packagesToQuery, fmt.Sprintf("//pkg/%s/...", info.Name()))
 	}
-	allPackages := strings.Join(packagesToQuery, "+")
+	packagesToQuery = append(packagesToQuery, "//pkg/ui:*")
+	return strings.Join(packagesToQuery, "+")
+}
+
+func generateTestSuites() {
 	cmd := exec.Command(
 		"bazel", "query",
-		fmt.Sprintf(`kind("(_get_x_data|(go|sh)_(binary|library|test|transition_binary|transition_test))", %s)`, allPackages),
+		fmt.Sprintf(`kind("(_get_x_data|(go|sh)_(binary|library|test|transition_binary|transition_test))", %s)`, getPackagesToQuery()),
 		"--output=label_kind",
 	)
 	buf, err := cmd.Output()
@@ -245,19 +285,19 @@ unused_checker(srcs = GET_X_DATA_TARGETS)`)
 }
 
 func generateTestsTimeouts() {
+	targets, err := getTestTargets()
+	if err != nil {
+		log.Fatal(err)
+	}
 	for size, timeout := range testSizeToDefaultTimeout {
-		targets, err := getTestTargets(size)
-		if err != nil {
-			log.Fatal(err)
-		}
 		// Let the `go test` process timeout 5 seconds before bazel attempts to kill it.
 		// Note that if this causes issues such as not having enough time to run normally
 		// (because of the 5 seconds taken) then the troubled test target size must be bumped
 		// to the next size because it shouldn't be passing at the edge of its deadline
 		// anyways to avoid flakiness.
 		runBuildozer(append([]string{
 			fmt.Sprintf(`set args "-test.timeout=%ds"`, timeout-5)},
-			targets...,
+			targets[size]...,
 		))
 	}
 }

pkg/cmd/generate-bazel-extra/main_test.go

@@ -0,0 +1,48 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseQueryXML(t *testing.T) {
+	expectedDataMap := map[string][]string{
+		"small": {
+			"//pkg/sql/sem/eval/cast_test:cast_test_test",
+			"//pkg/sql/sem/eval/eval_test:eval_test_test",
+		},
+		"medium": {
+			"//pkg/sql/sem/builtins:builtins_test",
+			"//pkg/sql/sem/builtins/pgformat:pgformat_test",
+			"//pkg/sql/sem/tree:tree_test",
+		},
+		"large": {
+			"//pkg/sql/sem/cast:cast_test",
+		},
+		"enormous": {
+			"//pkg/sql/sem/eval:eval_test",
+			"//pkg/sql/sem/normalize:normalize_test",
+		},
+	}
+	xmlData, err := os.ReadFile(filepath.Join(testutils.TestDataPath(t, "TestParseQueryXML"), "tc1.xml"))
+	require.NoError(t, err)
+	sizeToTargets, err := parseQueryXML(xmlData)
+	require.NoError(t, err)
+	for size := range expectedDataMap {
+		require.ElementsMatch(t, expectedDataMap[size], sizeToTargets[size])
+	}
+}