Merge cockroachdb#62116

62116: multiregionccl: unskip TestMultiRegionAfterEnterpriseDisabled r=ajstorm a=otan

Relevant fixes are now committed.

Resolves cockroachdb#61163 

Release note: None

Co-authored-by: Oliver Tan <otan@cockroachlabs.com>

pkg/ccl/multiregionccl/multiregion_test.go

@@ -16,7 +16,6 @@ import (
 	_ "github.com/cockroachdb/cockroach/pkg/ccl"
 	"github.com/cockroachdb/cockroach/pkg/ccl/multiregionccl/multiregionccltestutils"
 	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
@@ -54,22 +53,24 @@ func TestMultiRegionAfterEnterpriseDisabled(t *testing.T) {
 	defer log.Scope(t).Close(t)
 	defer utilccl.TestingEnableEnterprise()()
 
-	skip.UnderRace(t, "#61163")
-
 	_, sqlDB, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
 		t, 3 /* numServers */, base.TestingKnobs{}, nil, /* baseDir */
 	)
 	defer cleanup()
 
-	_, err := sqlDB.Exec(`
-CREATE DATABASE test PRIMARY REGION "us-east1" REGIONS "us-east2";
-USE test;
-CREATE TABLE t1 () LOCALITY GLOBAL;
-CREATE TABLE t2 () LOCALITY REGIONAL BY TABLE;
-CREATE TABLE t3 () LOCALITY REGIONAL BY TABLE IN "us-east2";
-CREATE TABLE t4 () LOCALITY REGIONAL BY ROW
-	`)
-	require.NoError(t, err)
+	for _, setupQuery := range []string{
+		`CREATE DATABASE test PRIMARY REGION "us-east1" REGIONS "us-east2"`,
+		`USE test`,
+		`CREATE TABLE t1 () LOCALITY GLOBAL`,
+		`CREATE TABLE t2 () LOCALITY REGIONAL BY TABLE`,
+		`CREATE TABLE t3 () LOCALITY REGIONAL BY TABLE IN "us-east2"`,
+		`CREATE TABLE t4 () LOCALITY REGIONAL BY ROW`,
+	} {
+		t.Run(setupQuery, func(t *testing.T) {
+			_, err := sqlDB.Exec(setupQuery)
+			require.NoError(t, err)
+		})
+	}
 
 	defer utilccl.TestingDisableEnterprise()()
 

pkg/ccl/utilccl/license_check.go

@@ -12,6 +12,7 @@ import (
 	"bytes"
 	"context"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -44,31 +45,41 @@ var enterpriseLicense = func() *settings.StringSetting {
 	return s
 }()
 
-var testingEnterpriseEnabled = false
+// testingEnterprise determines whether the cluster is enabled
+// or disabled for the purposes of testing.
+// It should be loaded and stored using atomic as it can race with an
+// in progress kv reader during TestingDisableEnterprise /
+// TestingEnableEnterprise.
+var testingEnterprise int32
+
+const (
+	testingEnterpriseDisabled = 0
+	testingEnterpriseEnabled  = 1
+)
 
 // TestingEnableEnterprise allows overriding the license check in tests.
 func TestingEnableEnterprise() func() {
-	before := testingEnterpriseEnabled
-	testingEnterpriseEnabled = true
+	before := atomic.LoadInt32(&testingEnterprise)
+	atomic.StoreInt32(&testingEnterprise, testingEnterpriseEnabled)
 	return func() {
-		testingEnterpriseEnabled = before
+		atomic.StoreInt32(&testingEnterprise, before)
 	}
 }
 
 // TestingDisableEnterprise allows re-enabling the license check in tests.
 func TestingDisableEnterprise() func() {
-	before := testingEnterpriseEnabled
-	testingEnterpriseEnabled = false
+	before := atomic.LoadInt32(&testingEnterprise)
+	atomic.StoreInt32(&testingEnterprise, testingEnterpriseDisabled)
 	return func() {
-		testingEnterpriseEnabled = before
+		atomic.StoreInt32(&testingEnterprise, before)
 	}
 }
 
 // CheckEnterpriseEnabled returns a non-nil error if the requested enterprise
 // feature is not enabled, including information or a link explaining how to
 // enable it.
 func CheckEnterpriseEnabled(st *cluster.Settings, cluster uuid.UUID, org, feature string) error {
-	if testingEnterpriseEnabled {
+	if atomic.LoadInt32(&testingEnterprise) == testingEnterpriseEnabled {
 		return nil
 	}
 	return checkEnterpriseEnabledAt(st, timeutil.Now(), cluster, org, feature)