sql: add tests for concurrent add/drop region operations

pkg/ccl/multiregionccl/region_test.go

@@ -24,91 +24,179 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestConcurrentDropRegion(t *testing.T) {
+// TestConcurrentAddDropRegions tests all combinations of add/drop as if they
+// were executed by two concurrent sessions. The general sketch of the test is
+// as follows:
+// - First operation is executed and blocks before the enum members are promoted.
+// - The second operation starts once the first operation has reached the type
+//   schema changer. It continues to completion. It may succeed/fail depending
+//   on the specific test setup.
+// - The first operation is resumed and allowed to complete. We expect it to
+//   succeed.
+// - Verify database regions are as expected.
+// Operations act on a multi-region database that contains a REGIONAL BY ROW
+// table, so as to exercise the repartitioning semantics.
+func TestConcurrentAddDropRegions(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	firstDropAtSchemaChange := false
-	secondDropAtSchemaChange := false
-	waitForFirstDropRegionToReachSchemaChanger := make(chan struct{})
-	waitForSecondDropRegionToReachSchemaChanger := make(chan struct{})
-	waitForFirstDropRegionToFinish := make(chan struct{})
-	var mu syncutil.Mutex
+	skip.UnderRace(t, "times out under race")
 
-	// Pause the first drop region in the type schema changer before it does enum
-	// member promotion. Then release it when the second drop region gets to the
-	// same state. This validates that the second drop region can complete without
-	// requiring an override (something which was not possible before #62354).
-	knobs := base.TestingKnobs{
-		SQLTypeSchemaChanger: &sql.TypeSchemaChangerTestingKnobs{
-			RunBeforeEnumMemberPromotion: func() {
-				mu.Lock()
-				if !firstDropAtSchemaChange {
-					firstDropAtSchemaChange = true
-					close(waitForFirstDropRegionToReachSchemaChanger)
-					mu.Unlock()
-					<-waitForSecondDropRegionToReachSchemaChanger
-				} else if !secondDropAtSchemaChange {
-					secondDropAtSchemaChange = true
-					// We're the second DROP REGION, close the channel to free the first
-					// DROP REGION.
-					close(waitForSecondDropRegionToReachSchemaChanger)
-					mu.Unlock()
-				}
-			},
+	testCases := []struct {
+		name            string
+		firstOp         string
+		secondOp        string
+		expectedRegions []string
+		secondOpErr     string
+	}{
+		{
+			"concurrent-drops",
+			`ALTER DATABASE db DROP REGION "us-east2"`,
+			`ALTER DATABASE db DROP REGION "us-east3"`,
+			[]string{"us-east1"},
+			"",
+		},
+		{
+			"concurrent-adds",
+			`ALTER DATABASE db ADD REGION "us-east4"`,
+			`ALTER DATABASE db ADD REGION "us-east5"`,
+			[]string{"us-east1", "us-east2", "us-east3", "us-east4", "us-east5"},
+			"",
+		},
+		{
+			"concurrent-add-drop",
+			`ALTER DATABASE db ADD REGION "us-east4"`,
+			`ALTER DATABASE db DROP REGION "us-east3"`,
+			[]string{"us-east1", "us-east2", "us-east4"},
+			"",
+		},
+		{
+			"concurrent-drop-add",
+			`ALTER DATABASE db DROP REGION "us-east2"`,
+			`ALTER DATABASE db ADD REGION "us-east5"`,
+			[]string{"us-east1", "us-east3", "us-east5"},
+			"",
+		},
+		{
+			"concurrent-add-same-region",
+			`ALTER DATABASE db ADD REGION "us-east5"`,
+			`ALTER DATABASE db ADD REGION "us-east5"`,
+			[]string{"us-east1", "us-east2", "us-east3", "us-east5"},
+			`region "us-east5" already added to database`,
+		},
+		{
+			"concurrent-drop-same-region",
+			`ALTER DATABASE db DROP REGION "us-east2"`,
+			`ALTER DATABASE db DROP REGION "us-east2"`,
+			[]string{"us-east1", "us-east3"},
+			`enum value "us-east2" is already being dropped`,
+		},
+		{
+			"concurrent-add-drop-same-region",
+			`ALTER DATABASE db ADD REGION "us-east5"`,
+			`ALTER DATABASE db DROP REGION "us-east5"`,
+			[]string{"us-east1", "us-east2", "us-east3", "us-east5"},
+			`enum value "us-east5" is being added, try again later`,
+		},
+		{
+			"concurrent-drop-add-same-region",
+			`ALTER DATABASE db DROP REGION "us-east2"`,
+			`ALTER DATABASE db ADD REGION "us-east2"`,
+			[]string{"us-east1", "us-east3"},
+			`enum value "us-east2" is being dropped, try again later`,
 		},
 	}
 
-	_, sqlDB, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
-		t, 3 /* numServers */, knobs, nil, /* baseDir */
-	)
-	defer cleanup()
-
-	_, err := sqlDB.Exec(`CREATE DATABASE db WITH PRIMARY REGION "us-east1" REGIONS "us-east2", "us-east3"`)
-	if err != nil {
-		t.Error(err)
-	}
-
-	// Send the first DROP REGION statement on its way.
-	go func() {
-		if _, err := sqlDB.Exec(`ALTER DATABASE db DROP REGION "us-east2"`); err != nil {
-			t.Error(err)
-		}
-		close(waitForFirstDropRegionToFinish)
-	}()
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			firstOp := true
+			firstOpStarted := make(chan struct{})
+			secondOpFinished := make(chan struct{})
+			firstOpFinished := make(chan struct{})
+			var mu syncutil.Mutex
+
+			knobs := base.TestingKnobs{
+				SQLTypeSchemaChanger: &sql.TypeSchemaChangerTestingKnobs{
+					RunBeforeEnumMemberPromotion: func() {
+						mu.Lock()
+						if firstOp {
+							firstOp = false
+							close(firstOpStarted)
+							mu.Unlock()
+							// Don't promote any members before the second operation reaches
+							// the schema changer as well.
+							<-secondOpFinished
+						} else {
+							mu.Unlock()
+						}
+					},
+				},
+			}
 
-	// Wait for the first DROP REGION to make it to the schema changer.
-	<-waitForFirstDropRegionToReachSchemaChanger
+			_, sqlDB, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
+				t, 5 /* numServers */, knobs, nil, /* baseDir */
+			)
+			defer cleanup()
 
-	// Issue the second DROP REGION statement.
-	_, err = sqlDB.Exec(`ALTER DATABASE db DROP REGION "us-east3"`)
+			// Create a multi-region database with a REGIONAL BY ROW table inside of it
+			// which needs to be re-partitioned on add/drop operations.
+			_, err := sqlDB.Exec(`
+CREATE DATABASE db WITH PRIMARY REGION "us-east1" REGIONS "us-east2", "us-east3"; 
+CREATE TABLE db.rbr () LOCALITY REGIONAL BY ROW`)
+			require.NoError(t, err)
 
-	if err != nil {
-		t.Error(err)
-	}
+			go func() {
+				if _, err := sqlDB.Exec(tc.firstOp); err != nil {
+					t.Error(err)
+				}
+				close(firstOpFinished)
+			}()
+
+			// Wait for the first operation to reach the type schema changer.
+			<-firstOpStarted
+
+			// Start the second operation.
+			_, err = sqlDB.Exec(tc.secondOp)
+			if tc.secondOpErr == "" {
+				require.NoError(t, err)
+			} else {
+				require.True(t, testutils.IsError(err, tc.secondOpErr))
+			}
+			close(secondOpFinished)
 
-	<-waitForFirstDropRegionToFinish
+			<-firstOpFinished
 
-	// Validate that both DROP REGION statements completed.
-	rows, err := sqlDB.Query("SELECT region FROM [SHOW REGIONS FROM DATABASE db]")
-	if err != nil {
-		t.Error(err)
-	}
-	defer rows.Close()
+			dbRegions := make([]string, 0, len(tc.expectedRegions))
+			rows, err := sqlDB.Query("SELECT region FROM [SHOW REGIONS FROM DATABASE db]")
+			require.NoError(t, err)
+			for {
+				done := rows.Next()
+				if !done {
+					break
+				}
+				var region string
+				err := rows.Scan(&region)
+				require.NoError(t, err)
 
-	const expectedRegion = "us-east1"
-	var region string
-	rows.Next()
-	if err := rows.Scan(&region); err != nil {
-		t.Error(err)
-	}
+				dbRegions = append(dbRegions, region)
+			}
 
-	if region != expectedRegion {
-		t.Error(errors.Newf("expected region to be: %q, got %q", expectedRegion, region))
-	}
+			if len(dbRegions) != len(tc.expectedRegions) {
+				t.Fatalf("unexpected number of regions, expected: %v found %v",
+					tc.expectedRegions,
+					dbRegions,
+				)
+			}
 
-	if rows.Next() {
-		t.Error(errors.New("unexpected number of rows returned"))
+			for i, expectedRegion := range tc.expectedRegions {
+				if expectedRegion != dbRegions[i] {
+					t.Fatalf("unexpected regions, expected: %v found %v",
+						tc.expectedRegions,
+						dbRegions,
+					)
+				}
+			}
+		})
 	}
 }
 
@@ -164,7 +252,7 @@ ALTER DATABASE db PRIMARY REGION "us-east2";
 		t.Fatalf("expected error, found nil")
 	}
 	if !testutils.IsError(err, `"us-east2" has not been added to the database`) {
-		t.Fatalf(`expected err, got %v`, err)
+		t.Fatalf(`expected secondOpErr, got %v`, err)
 	}
 
 	close(waitForPrimaryRegionSwitch)

pkg/sql/type_change.go

@@ -459,11 +459,6 @@ func performMultiRegionFinalization(
 // all REGIONAL BY ROW tables in the enclosing database such that there is a
 // partition and corresponding zone configuration for all PUBLIC enum members
 // (regions).
-// This currently doesn't work too well if there are READ ONLY member on the
-// type. This is because we create the partitioning clause based on the regions
-// on the database descriptor, which may include the READ ONLY member, and
-// partitioning on a READ ONLY member doesn't work. This will go away once
-// https://github.com/cockroachdb/cockroach/issues/60620 is fixed.
 func repartitionRegionalByRowTables(
 	ctx context.Context,
 	typeDesc *typedesc.Immutable,