Revert "bench/rttanalysis: shard TestBenchmarkExpectation to avoid timeouts"

This reverts commit 9fecc53.

Fixes: #154317
Release note: none

Author: fqazi

pkg/bench/rttanalysis/BUILD.bazel

@@ -13,8 +13,6 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/base",
-        "//pkg/jobs",
-        "//pkg/jobs/jobspb",
         "//pkg/kv/kvclient/kvcoord",
         "//pkg/sql",
         "//pkg/sql/parser",
@@ -58,9 +56,9 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":rttanalysis"],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 4,
     deps = [
         "//pkg/base",
+        "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
@@ -71,7 +69,6 @@ go_test(
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/testcluster",
-        "//pkg/util/envutil",
         "//pkg/util/protoutil",
         "//pkg/util/randutil",
     ],

pkg/bench/rttanalysis/registry.go

@@ -9,8 +9,6 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/cockroachdb/cockroach/pkg/jobs"
-	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
@@ -53,66 +51,15 @@ 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 duress and short.
+// It takes a long time and thus is skipped under stress, race
+// and short.
 func (r *Registry) RunExpectations(t *testing.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) {
-	defer jobs.TestingSetIDsToIgnore(map[jobspb.JobID]struct{}{3001: {}, 3002: {}})()
-	skip.UnderDuress(t)
+	skip.UnderStress(t)
+	skip.UnderRace(t)
 	skip.UnderShort(t)
+	skip.UnderDeadlock(t)
 
-	// 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, registryToUse)
+	runBenchmarkExpectationTests(t, r)
 }
 
 // Register registers a set of test cases to a given benchmark name. It is

pkg/bench/rttanalysis/validate_benchmark_data_test.go

@@ -6,44 +6,13 @@
 package rttanalysis
 
 import (
-	"strconv"
 	"testing"
 
-	"github.com/cockroachdb/cockroach/pkg/util/envutil"
+	"github.com/cockroachdb/cockroach/pkg/jobs"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 )
 
-// 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)
+func TestBenchmarkExpectation(t *testing.T) {
+	defer jobs.TestingSetIDsToIgnore(map[jobspb.JobID]struct{}{3001: {}, 3002: {}})()
+	reg.RunExpectations(t)
 }