From bf9a6a4ee9af22c98e88cb8976235350ae97292c Mon Sep 17 00:00:00 2001
From: arulajmani <arulajmani@gmail.com>
Date: Wed, 10 Feb 2021 22:58:34 -0500
Subject: [PATCH] sql: test namespace entry for multi-region enum is drained
 properly

Testing only patch. When the primary region is removed from a
multi-region database we queue up an async job to reclaim the namespace
entry for the multi-region type descriptor (and its array alias). This
patch adds a test to make sure that the namespace entry is drained
properly even if the job runs into an error.

This patch adds a new testing knob to the type schema changer,
RunAfterAllFailOrCancel, to synch on the `OnFailOrCancel` job
completing.

Release note: None
---
 pkg/sql/multiregion_test.go | 77 +++++++++++++++++++++++++++++++++++++
 pkg/sql/type_change.go      | 15 +++++++-
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/pkg/sql/multiregion_test.go b/pkg/sql/multiregion_test.go
index bc4a20aabbee..ef8025b60f5e 100644
--- a/pkg/sql/multiregion_test.go
+++ b/pkg/sql/multiregion_test.go
@@ -18,12 +18,15 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqltestutils"
+	"github.com/cockroachdb/cockroach/pkg/sql/tests"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSettingPrimaryRegionAmidstDrop(t *testing.T) {
@@ -127,3 +130,77 @@ func TestSettingPrimaryRegionAmidstDrop(t *testing.T) {
 		return nil
 	})
 }
+
+// TestDroppingPrimaryRegionAsyncJobFailure drops the primary region of the
+// database, which results in dropping the multi-region type descriptor. Then,
+// it errors out the async job associated with the type descriptor cleanup and
+// ensures the namespace entry is reclaimed back despite the injected error.
+// We rely on this behavior to be able to add multi-region capability in the
+// future.
+func TestDroppingPrimaryRegionAsyncJobFailure(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// Decrease the adopt loop interval so that retries happen quickly.
+	defer sqltestutils.SetTestJobsAdoptInterval()()
+
+	params, _ := tests.CreateTestServerParams()
+	params.Locality.Tiers = []roachpb.Tier{
+		{Key: "region", Value: "us-east-1"},
+	}
+
+	// Protects expectedCleanupRuns
+	var mu syncutil.Mutex
+	// We need to cleanup 2 times, once for the multi-region type descriptor and
+	// once for the array alias of the multi-region type descriptor.
+	expectedCleanupRuns := 2
+	cleanupFinished := make(chan struct{})
+	params.Knobs.SQLTypeSchemaChanger = &sql.TypeSchemaChangerTestingKnobs{
+		RunBeforeExec: func() error {
+			return errors.New("yikes")
+		},
+		RunAfterOnFailOrCancel: func() error {
+			mu.Lock()
+			defer mu.Unlock()
+			expectedCleanupRuns--
+			if expectedCleanupRuns == 0 {
+				close(cleanupFinished)
+			}
+			return nil
+		},
+	}
+
+	s, sqlDB, _ := serverutils.StartServer(t, params)
+	ctx := context.Background()
+	defer s.Stopper().Stop(ctx)
+
+	// Setup the test.
+	_, err := sqlDB.Exec(`
+CREATE DATABASE db WITH PRIMARY REGION "us-east-1"; 
+CREATE TABLE db.t(k INT) LOCALITY REGIONAL BY TABLE IN PRIMARY REGION;
+`)
+	require.NoError(t, err)
+
+	_, err = sqlDB.Exec(`ALTER DATABASE db DROP REGION "us-east-1"`)
+	testutils.IsError(err, "yikes")
+
+	<-cleanupFinished
+
+	rows := sqlDB.QueryRow(`SELECT COUNT(*) FROM system.namespace WHERE name = 'crdb_internal_region'`)
+	var count int
+	err = rows.Scan(&count)
+	require.NoError(t, err)
+	if count != 0 {
+		t.Fatal("expected crdb_internal_region not to be present in system.namespace")
+	}
+
+	_, err = sqlDB.Exec(`ALTER DATABASE db PRIMARY REGION "us-east-1"`)
+	require.NoError(t, err)
+
+	rows = sqlDB.QueryRow(`SELECT COUNT(*) FROM system.namespace WHERE name = 'crdb_internal_region'`)
+	err = rows.Scan(&count)
+	require.NoError(t, err)
+	if count != 1 {
+		t.Fatal("expected crdb_internal_region to be present in system.namespace")
+	}
+}
diff --git a/pkg/sql/type_change.go b/pkg/sql/type_change.go
index 2f9c8490a5ec..29a1eb387c79 100644
--- a/pkg/sql/type_change.go
+++ b/pkg/sql/type_change.go
@@ -162,6 +162,9 @@ type TypeSchemaChangerTestingKnobs struct {
 	// RunBeforeEnumMemberPromotion runs before enum members are promoted from
 	// readable to all permissions in the typeSchemaChanger.
 	RunBeforeEnumMemberPromotion func()
+	// RunAfterOnFailOrCancel runs after OnFailOrCancel completes, if
+	// OnFailOrCancel is triggered.
+	RunAfterOnFailOrCancel func() error
 }
 
 // ModuleTestingKnobs implements the ModuleTestingKnobs interface.
@@ -616,10 +619,18 @@ func (t *typeChangeResumer) OnFailOrCancel(ctx context.Context, execCtx interfac
 		return err
 	}
 
-	return drainNamesForDescriptor(
+	if err := drainNamesForDescriptor(
 		ctx, tc.execCfg.Settings, tc.typeID, tc.execCfg.DB,
 		tc.execCfg.InternalExecutor, tc.execCfg.LeaseManager, tc.execCfg.Codec, nil,
-	)
+	); err != nil {
+		return err
+	}
+
+	if fn := tc.execCfg.TypeSchemaChangerTestingKnobs.RunAfterOnFailOrCancel; fn != nil {
+		return fn()
+	}
+
+	return nil
 }
 
 func init() {