release-24.3: bench/rttanalysis: shard TestBenchmarkExpectation to avoid timeouts

pkg/bench/rttanalysis/BUILD.bazel

@@ -55,6 +55,7 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":rttanalysis"],
     exec_properties = {"test.Pool": "large"},
+    shard_count = 4,
     deps = [
         "//pkg/base",
         "//pkg/jobs/jobspb",
@@ -66,6 +67,7 @@ go_test(
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
+        "//pkg/util/envutil",
         "//pkg/util/protoutil",
         "//pkg/util/randutil",
     ],

pkg/bench/rttanalysis/registry.go

@@ -6,6 +6,7 @@
 package rttanalysis
 
 import (
+	"runtime"
 	"strings"
 	"testing"
 
@@ -51,15 +52,68 @@ func (r *Registry) Run(b *testing.B) {
 // benchmarks can be filtered by passing the usual test filters underneath
 // this test's name.
 //
-// It takes a long time and thus is skipped under stress, race
-// and short.
+// It takes a long time and thus is skipped under duress and short.
 func (r *Registry) RunExpectations(t *testing.T) {
-	skip.UnderStress(t)
-	skip.UnderRace(t)
+	r.RunExpectationsSharded(t, 1, 1)
+}
+
+// RunExpectationsSharded runs all the benchmarks for one iteration
+// and validates that the number of RPCs meets the expectation. If run
+// with the --rewrite flag, it will rewrite the run benchmarks. The
+// benchmarks can be filtered by passing the usual test filters underneath
+// this test's name.
+//
+// It takes a long time and thus is skipped under duress and short.
+//
+// When shard and totalShards are provided (> 1), only a subset of benchmarks
+// assigned to the specific shard will be run, enabling parallel execution.
+// Test groups are distributed across shards using round-robin assignment.
+func (r *Registry) RunExpectationsSharded(t *testing.T, shard, totalShards int) {
+	skip.UnderDuress(t)
 	skip.UnderShort(t)
-	skip.UnderDeadlock(t)
+	if runtime.GOARCH == "s390x" {
+		skip.IgnoreLint(t, "test prone to crashing under s390x (see #154317)")
+	}
+
+	// If totalShards is 1, run all tests; otherwise shard them
+	var registryToUse *Registry
+	if totalShards <= 1 {
+		// Run all test groups
+		registryToUse = r
+	} else {
+		// Create a registry with only the test groups assigned to this shard
+		shardRegistry := &Registry{
+			numNodes: r.numNodes,
+			cc:       r.cc,
+			r:        make(map[string][]RoundTripBenchTestCase),
+		}
+
+		// Distribute test groups across shards using round-robin assignment
+		// First, get all group names and sort them for consistent ordering
+		groupNames := make([]string, 0, len(r.r))
+		for groupName := range r.r {
+			groupNames = append(groupNames, groupName)
+		}
+		// Sort for deterministic assignment across runs
+		for i := 0; i < len(groupNames); i++ {
+			for j := i + 1; j < len(groupNames); j++ {
+				if groupNames[i] > groupNames[j] {
+					groupNames[i], groupNames[j] = groupNames[j], groupNames[i]
+				}
+			}
+		}
+
+		// Assign groups to shards using round-robin
+		for i, groupName := range groupNames {
+			assignedShard := (i % totalShards) + 1
+			if assignedShard == shard {
+				shardRegistry.r[groupName] = r.r[groupName]
+			}
+		}
+		registryToUse = shardRegistry
+	}
 
-	runBenchmarkExpectationTests(t, r)
+	runBenchmarkExpectationTests(t, registryToUse)
 }
 
 // Register registers a set of test cases to a given benchmark name. It is

pkg/bench/rttanalysis/validate_benchmark_data_test.go

@@ -5,6 +5,45 @@
 
 package rttanalysis
 
-import "testing"
+import (
+	"strconv"
+	"testing"
 
-func TestBenchmarkExpectation(t *testing.T) { reg.RunExpectations(t) }
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
+)
+
+// NOTE: If you change the number of shards, you must also update the
+// shard_count in BUILD.bazel to match.
+const shardCount = 4
+
+// Validate that shardCount matches TEST_TOTAL_SHARDS environment variable at init time
+var _ = func() int {
+	totalShardsStr, found := envutil.ExternalEnvString("TEST_TOTAL_SHARDS", 1)
+	if totalShardsStr == "" || !found {
+		return 0
+	}
+	totalShards, err := strconv.Atoi(totalShardsStr)
+	if err != nil {
+		return 0
+	}
+	if totalShards != shardCount {
+		panic("shardCount mismatch: update shard_count in pkg/bench/rttanalysis/BUILD.bazel to match shardCount constant")
+	}
+	return 0
+}()
+
+func TestBenchmarkExpectationShard1(t *testing.T) {
+	reg.RunExpectationsSharded(t, 1, shardCount)
+}
+
+func TestBenchmarkExpectationShard2(t *testing.T) {
+	reg.RunExpectationsSharded(t, 2, shardCount)
+}
+
+func TestBenchmarkExpectationShard3(t *testing.T) {
+	reg.RunExpectationsSharded(t, 3, shardCount)
+}
+
+func TestBenchmarkExpectationShard4(t *testing.T) {
+	reg.RunExpectationsSharded(t, 4, shardCount)
+}