drpc: add DefaultTestDRPCEnableOption to TestServerArgs

This entire setup is similar to `DefaultTestTenantOptions`.

This option can be configured to run a server with DRPC enabled. The plan
is to enable as many test packages as possible at the global level using
`TestingGlobalDRPCEnableOption` in `TestMain` and, for a select few others,
just specify this option individually in their `TestServerArgs`.

One follow-up to this PR is to simulate a mixed-mode server where, in
`TestClusterArgs.ServerArgsPerNode`, we specify some nodes with DRPC
enabled and others not. This would require some testing knobs because
enabling a setting on one node will take effect for others as well, and we
would require some refactoring in other production code paths too.

Release note: none
Epic: CRDB-48935

Author: shubhamdhama

pkg/base/test_server_args.go

@@ -159,6 +159,11 @@ type TestServerArgs struct {
 	// below for alternative options that suits your test case.
 	DefaultTestTenant DefaultTestTenantOptions
 
+	// DefaultDRPCOption specifies the DRPC enablement mode for a test
+	// server. This controls whether inter-node connectivity uses DRPC, just
+	// gRPC, or is chosen randomly.
+	DefaultDRPCOption DefaultTestDRPCOption
+
 	// DefaultTenantName is the name of the tenant created implicitly according
 	// to DefaultTestTenant. It is typically `test-tenant` for unit tests and
 	// always `demoapp` for the cockroach demo.
@@ -224,6 +229,25 @@ type SlimTestServerConfig struct {
 	Options slimOptions
 }
 
+// DefaultTestDRPCOption specifies the DRPC enablement mode for a test
+// server. This controls whether inter-node connectivity uses DRPC, just gRPC,
+// or is chosen randomly.
+type DefaultTestDRPCOption uint8
+
+const (
+	// TestDRPCDisabled disables DRPC; all inter-node connectivity will use gRPC
+	// only.
+	TestDRPCDisabled DefaultTestDRPCOption = iota
+
+	// TestDRPCEnabled enables DRPC. Some services may still use gRPC if they
+	// have not yet migrated to DRPC.
+	TestDRPCEnabled
+
+	// TestDRPCEnabledRandomly randomly chooses between the behavior of
+	// TestDRPCDisabled or TestDRPCEnabled.
+	TestDRPCEnabledRandomly
+)
+
 // 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/drpc_test.go

@@ -15,6 +15,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -159,3 +161,21 @@ func TestStreamContextCancel(t *testing.T) {
 		singleRequest()
 	}
 }
+
+func TestDefaultDRPCOption(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	testutils.RunTrueAndFalse(t, "drpc-enabled", func(t *testing.T, option bool) {
+		opt := base.TestDRPCDisabled
+		if option {
+			opt = base.TestDRPCEnabled
+		}
+		s := serverutils.StartServerOnly(t, base.TestServerArgs{DefaultDRPCOption: opt})
+		defer s.Stopper().Stop(context.Background())
+
+		var enabled bool
+		sqlutils.MakeSQLRunner(s.SQLConn(t)).QueryRow(t, "SHOW CLUSTER SETTING rpc.experimental_drpc.enabled").Scan(&enabled)
+		require.Equal(t, option, enabled)
+	})
+}

pkg/testutils/serverutils/test_server_shim.go

@@ -24,8 +24,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilitiespb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils/pgurlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
@@ -203,6 +205,23 @@ func testTenantDecisionFromEnvironment(
 	return baseArg, false
 }
 
+var globalDefaultDRPCOptionOverride struct {
+	isSet bool
+	value base.DefaultTestDRPCOption
+}
+
+// TestingGlobalDRPCOption sets the package-level DefaultTestDRPCOption.
+//
+// Note: This override will be superseded by any more specific options provided
+// when starting the server or cluster.
+func TestingGlobalDRPCOption(v base.DefaultTestDRPCOption) func() {
+	globalDefaultDRPCOptionOverride.isSet = true
+	globalDefaultDRPCOptionOverride.value = v
+	return func() {
+		globalDefaultDRPCOptionOverride.isSet = false
+	}
+}
+
 // globalDefaultSelectionOverride is used when an entire package needs
 // to override the probabilistic behavior.
 var globalDefaultSelectionOverride struct {
@@ -252,11 +271,15 @@ type TestFataler interface {
 // The first argument is optional. If non-nil; it is used for logging
 // server configuration messages.
 func StartServerOnlyE(t TestLogger, params base.TestServerArgs) (TestServerInterface, error) {
+	ctx := context.Background()
 	allowAdditionalTenants := params.DefaultTestTenant.AllowAdditionalTenants()
+
 	// Update the flags with the actual decision as to whether we should
 	// start the service for a default test tenant.
 	params.DefaultTestTenant = ShouldStartDefaultTestTenant(t, params.DefaultTestTenant)
 
+	TryEnableDRPCSetting(ctx, t, &params)
+
 	s, err := NewServer(params)
 	if err != nil {
 		return nil, err
@@ -269,8 +292,6 @@ func StartServerOnlyE(t TestLogger, params base.TestServerArgs) (TestServerInter
 		}
 	}
 
-	ctx := context.Background()
-
 	if err := s.Start(ctx); err != nil {
 		return nil, err
 	}
@@ -519,3 +540,29 @@ func WaitForTenantCapabilities(
 		t.Fatal(err)
 	}
 }
+
+// TryEnableDRPCSetting determines whether to enable the DRPC cluster setting
+// based on the `TestServerArgs.DefaultDRPCOption` and updates the
+// `TestServerArgs.Settings` based on that.
+func TryEnableDRPCSetting(ctx context.Context, t TestLogger, args *base.TestServerArgs) {
+	option := args.DefaultDRPCOption
+	if option == base.TestDRPCDisabled && globalDefaultDRPCOptionOverride.isSet {
+		option = globalDefaultDRPCOptionOverride.value
+	}
+	enableDRPC := false
+	switch option {
+	case base.TestDRPCEnabled:
+		enableDRPC = true
+	case base.TestDRPCEnabledRandomly:
+		rng, _ := randutil.NewTestRand()
+		enableDRPC = rng.Intn(2) == 0
+	}
+	if !enableDRPC {
+		return
+	}
+
+	if args.Settings == nil {
+		args.Settings = cluster.MakeClusterSettings()
+	}
+	rpc.ExperimentalDRPCEnabled.Override(ctx, &args.Settings.SV, true)
+}

pkg/testutils/testcluster/testcluster.go

@@ -346,6 +346,8 @@ func NewTestCluster(
 			serverArgs.Settings = cluster.TestingCloneClusterSettings(serverArgs.Settings)
 		}
 
+		serverutils.TryEnableDRPCSetting(context.Background(), t, &serverArgs)
+
 		// If a reusable listener registry is provided, create reusable listeners
 		// for every server that doesn't have a custom listener provided. (Only
 		// servers with a reusable listener can be restarted).