Merge #154156

154156: testutils,drpc: improve experience with DefaultTestDRPCOption r=shubhamdhama a=shubhamdhama

Here are some small code changes to improve developer experience.

`` testutils: improve validate DefaultDRPCOption error message ``

This improves the error message when node level server args are used
instead of top-level server args for a test cluster.

before: "improper use of DefaultDRPCOption in per-server args: 1 vs 0"
now   : "improper use of DefaultDRPCOption in per-server args: enabled vs unset"

`` server,tests: explicitly disable DRPC setting if option is TestDRPCDisabled ``

This change ensures a clear distinction between an unset DRPC option and an
explicitly disabled one. If the DRPC option is unset at both the test and
package levels, we will not modify the setting. This allows the system to
use the default value of `ExperimentalDRPCEnabled`, which can be overridden
for development via `COCKROACH_EXPERIMENTAL_DRPC_ENABLED`. When the option
is explicitly set to `TestDRPCDisabled`, we will now explicitly disable the
setting.

`` testutils: allow overriding DefaultTestDRPCOption with env variable ``

This allows enabling or disabling DRPC for unit tests without changing code
or recompiling. Note that the environment variable will take precedence
over server args, which take precedence over the package-level option.

Example usage: `dev test pkg/server:server_test -- --test_env=COCKROACH_TEST_DRPC=false`

Epic: none
Release note: none

Co-authored-by: Shubham Dhama <[email protected]>

Author: craig[bot]

pkg/base/test_server_args.go

@@ -251,6 +251,21 @@ const (
 	TestDRPCEnabledRandomly
 )
 
+func (d DefaultTestDRPCOption) String() string {
+	switch d {
+	case TestDRPCUnset:
+		return "unset"
+	case TestDRPCDisabled:
+		return "disabled"
+	case TestDRPCEnabled:
+		return "enabled"
+	case TestDRPCEnabledRandomly:
+		return "enabled-randomly"
+	default:
+		panic("unreachable")
+	}
+}
+
 // TestClusterArgs contains the parameters one can set when creating a test
 // cluster. It contains a TestServerArgs instance which will be copied over to
 // every server.

pkg/server/testserver.go

@@ -333,8 +333,11 @@ func makeTestConfigFromParams(params base.TestServerArgs) Config {
 		cfg.TestingKnobs.AdmissionControlOptions = &admission.Options{}
 	}
 
-	if params.DefaultDRPCOption == base.TestDRPCEnabled {
+	switch params.DefaultDRPCOption {
+	case base.TestDRPCEnabled:
 		rpcbase.ExperimentalDRPCEnabled.Override(context.Background(), &st.SV, true)
+	case base.TestDRPCDisabled:
+		rpcbase.ExperimentalDRPCEnabled.Override(context.Background(), &st.SV, false)
 	}
 
 	return cfg

pkg/testutils/serverutils/test_server_shim.go

@@ -177,6 +177,13 @@ const (
 
 	testTenantModeEnabledShared   = "shared"
 	testTenantModeEnabledExternal = "external"
+
+	// COCKROACH_TEST_DRPC controls the DRPC enablement mode for test servers.
+	//
+	// - disabled: disables DRPC; all inter-node connectivity will use gRPC only
+	//
+	// - enabled: enables DRPC for inter-node connectivity
+	testDRPCEnabledEnvVar = "COCKROACH_TEST_DRPC"
 )
 
 func testTenantDecisionFromEnvironment(
@@ -537,24 +544,50 @@ func WaitForTenantCapabilities(
 	}
 }
 
+// parseDefaultTestDRPCOptionFromEnv parses the COCKROACH_TEST_DRPC environment
+// variable and returns the corresponding DefaultTestDRPCOption. If the
+// environment variable is not set it returns TestDRPCUnset. For invalid value,
+// it panic.
+func parseDefaultTestDRPCOptionFromEnv() base.DefaultTestDRPCOption {
+	if str, present := envutil.EnvString(testDRPCEnabledEnvVar, 0); present {
+		switch str {
+		case "disabled", "false":
+			return base.TestDRPCDisabled
+		case "enabled", "true":
+			return base.TestDRPCEnabled
+		default:
+			panic(fmt.Sprintf("invalid value for %s: %s", testDRPCEnabledEnvVar, str))
+		}
+	}
+	return base.TestDRPCUnset
+}
+
 // ShouldEnableDRPC determines the final DRPC option based on the input
 // option and any global overrides, resolving random choices to a concrete
 // enabled/disabled state.
 func ShouldEnableDRPC(
 	ctx context.Context, t TestLogger, option base.DefaultTestDRPCOption,
 ) base.DefaultTestDRPCOption {
 	var logSuffix string
-	if option == base.TestDRPCUnset && globalDefaultDRPCOptionOverride.isSet {
+
+	// Check environment variable first
+	if envOption := parseDefaultTestDRPCOptionFromEnv(); envOption != base.TestDRPCUnset {
+		option = envOption
+		logSuffix = " (override by COCKROACH_TEST_DRPC environment variable)"
+	} else if option == base.TestDRPCUnset && globalDefaultDRPCOptionOverride.isSet {
 		option = globalDefaultDRPCOptionOverride.value
 		logSuffix = " (override by TestingGlobalDRPCOption)"
 	}
+
 	enableDRPC := false
 	switch option {
 	case base.TestDRPCEnabled:
 		enableDRPC = true
 	case base.TestDRPCEnabledRandomly:
 		rng, _ := randutil.NewTestRand()
 		enableDRPC = rng.Intn(2) == 0
+	case base.TestDRPCUnset:
+		return base.TestDRPCUnset
 	}
 
 	if enableDRPC {

pkg/testutils/testcluster/testcluster.go

@@ -281,16 +281,16 @@ func (tc *TestCluster) validateDefaultTestTenant(
 // validateDefaultDRPCOption checks that per-server args don't override
 // the top-level DefaultDRPCOption setting inconsistently.
 func (tc *TestCluster) validateDefaultDRPCOption(
-	t serverutils.TestFataler, nodes int, defaultDRPCOption base.DefaultTestDRPCOption,
+	t serverutils.TestFataler, nodes int, clusterDRPCOption base.DefaultTestDRPCOption,
 ) {
 	for i := range nodes {
 		if args, ok := tc.clusterArgs.ServerArgsPerNode[i]; ok &&
 			args.DefaultDRPCOption != base.TestDRPCUnset &&
-			args.DefaultDRPCOption != defaultDRPCOption {
+			args.DefaultDRPCOption != clusterDRPCOption {
 			tc.Stopper().Stop(context.Background())
 			t.Fatalf("improper use of DefaultDRPCOption in per-server args: %v vs %v\n"+
-				"Tip: use the top-level ServerArgs to set the default DRPC option.",
-				args.DefaultDRPCOption, defaultDRPCOption)
+				"Use the top-level ServerArgs to set the default DRPC option.",
+				args.DefaultDRPCOption, clusterDRPCOption)
 		}
 	}
 }