Merge #148133 #148180

148133: drpc: add DefaultTestDRPCEnableOption to TestServerArgs r=cthumuluru-crdb a=shubhamdhama

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

148180: workload/schemachanger: handle foreign key mutations r=fqazi a=fqazi

Previously, the insert operation could fail on foreign key violation errors, if the foreign key constraint was enforce using a mutation. This patch detects those scenarios and turns these errors into potential errors.

Fixes: #148070
Fixes: #147741 
Fixes:  #147900

Release note: None

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

Author: 3 people

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).

pkg/workload/schemachange/error_screening.go

@@ -1726,3 +1726,32 @@ func (og *operationGenerator) tableHasUniqueConstraintMutation(
 			AND (m->'index'->>'unique')::BOOL IS TRUE
 		);`, tableName)
 }
+
+// tableHasForeignKeyMutation determines if a table has any foreign key constraint
+// mutation ongoing. This means either being added or dropped.
+func (og *operationGenerator) tableHasForeignKeyMutation(
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName,
+) (bool, error) {
+	return og.scanBool(ctx, tx, `
+		WITH table_desc AS (
+			SELECT crdb_internal.pb_to_json(
+				'desc',
+				descriptor,
+				false
+			)->'table' as d
+			FROM system.descriptor
+			WHERE id = $1::REGCLASS
+		)
+		SELECT EXISTS (
+			SELECT * FROM (
+			SELECT jsonb_array_elements(
+				CASE WHEN d->'mutations' IS NULL
+				THEN '[]'::JSONB
+				ELSE d->'mutations'
+				END
+			) as m
+			FROM table_desc)
+			WHERE (m->>'direction')::STRING IN ('ADD', 'DROP')
+			AND (m->'constraint'->>'foreign_key') IS NOT NULL
+		);`, tableName)
+}

pkg/workload/schemachange/operation_generator.go

@@ -3082,6 +3082,7 @@ func (og *operationGenerator) insertRow(ctx context.Context, tx pgx.Tx) (stmt *o
 	// we will need to evaluate generated expressions below.
 	hasUniqueConstraints := false
 	hasUniqueConstraintsMutations := false
+	hasForeignKeyMutations := false
 	fkViolation := false
 	if !anyInvalidInserts {
 		// Verify if the new row may violate unique constraints by checking the
@@ -3103,18 +3104,25 @@ func (og *operationGenerator) insertRow(ctx context.Context, tx pgx.Tx) (stmt *o
 		if err != nil {
 			return nil, err
 		}
-
+		// Determine if a foreign key mutation exists.
+		hasForeignKeyMutations, err = og.tableHasForeignKeyMutation(ctx, tx, tableName)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	stmt.potentialExecErrors.addAll(codesWithConditions{
 		{code: pgcode.UniqueViolation, condition: hasUniqueConstraints || hasUniqueConstraintsMutations},
-		{code: pgcode.ForeignKeyViolation, condition: fkViolation},
+		{code: pgcode.ForeignKeyViolation, condition: fkViolation || hasForeignKeyMutations},
 		{code: pgcode.NotNullViolation, condition: true},
 		{code: pgcode.CheckViolation, condition: true},
 		{code: pgcode.InsufficientPrivilege, condition: true}, // For RLS violations
 	})
 	og.expectedCommitErrors.addAll(codesWithConditions{
-		{code: pgcode.ForeignKeyViolation, condition: fkViolation},
+		{code: pgcode.ForeignKeyViolation, condition: fkViolation && !hasForeignKeyMutations},
+	})
+	og.potentialCommitErrors.addAll(codesWithConditions{
+		{code: pgcode.ForeignKeyViolation, condition: hasForeignKeyMutations},
 	})
 
 	var formattedRows []string