instancestorage: load all regions from the system DB when reclaiming IDs

Follow up to cockroachdb#92010 and cockroachdb#91694.

This commit addresses a TODO comment by loading all region enum members from
the system database when reclaiming IDs. The list of regions is read each time
the loop runs since there's a possibiilty where regions are added or removed
from the system DB.

Epic: None

Release note: None

pkg/ccl/multiregionccl/BUILD.bazel

@@ -64,6 +64,7 @@ go_test(
         "//pkg/settings/cluster",
         "//pkg/sql",
         "//pkg/sql/catalog",
+        "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/desctestutils",

pkg/ccl/multiregionccl/region_util_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -43,47 +44,80 @@ func TestGetLocalityRegionEnumPhysicalRepresentation(t *testing.T) {
 
 	s0 := tc.ServerTyped(0)
 	ief := s0.InternalExecutorFactory().(descs.TxnManager)
-	dbID := sqlutils.QueryDatabaseID(t, sqlDB, "foo")
+	dbID := descpb.ID(sqlutils.QueryDatabaseID(t, sqlDB, "foo"))
 
 	t.Run("with locality that exists", func(t *testing.T) {
 		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
-			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{
+			ctx, ief, s0.DB(), dbID, roachpb.Locality{
 				Tiers: []roachpb.Tier{{Key: "region", Value: "us-east2"}},
 			},
 		)
 		require.NoError(t, err)
 
-		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), dbID)
 		require.NotEmpty(t, enumMembers)
 		require.Equal(t, enumMembers["us-east2"], regionEnum)
 	})
 
 	t.Run("with non-existent locality", func(t *testing.T) {
 		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
-			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{
+			ctx, ief, s0.DB(), dbID, roachpb.Locality{
 				Tiers: []roachpb.Tier{{Key: "region", Value: "europe-west1"}},
 			},
 		)
 		require.NoError(t, err)
 
 		// Fallback to primary region if the locality is provided, but non-existent.
-		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), dbID)
 		require.NotEmpty(t, enumMembers)
 		require.Equal(t, enumMembers["us-east1"], regionEnum)
 	})
 
 	t.Run("without locality", func(t *testing.T) {
 		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
-			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{})
+			ctx, ief, s0.DB(), dbID, roachpb.Locality{})
 		require.NoError(t, err)
 
 		// Fallback to primary region is locality information is missing.
-		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), dbID)
 		require.NotEmpty(t, enumMembers)
 		require.Equal(t, enumMembers["us-east1"], regionEnum)
 	})
 }
 
+func TestGetRegionEnumRepresentations(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	tc, sqlDB, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
+		t, 3 /* numServers */, base.TestingKnobs{})
+	defer cleanup()
+
+	tDB := sqlutils.MakeSQLRunner(sqlDB)
+	tDB.Exec(t, `CREATE DATABASE foo PRIMARY REGION "us-east1" REGIONS "us-east1", "us-east2", "us-east3"`)
+
+	dbID := descpb.ID(sqlutils.QueryDatabaseID(t, sqlDB, "foo"))
+	err := sql.TestingDescsTxn(ctx, tc.Server(0), func(ctx context.Context, txn *kv.Txn, col *descs.Collection) error {
+		enumReps, primaryRegion, err := sql.GetRegionEnumRepresentations(ctx, txn, dbID, col)
+		require.NoError(t, err)
+
+		require.Equal(t, catpb.RegionName("us-east1"), primaryRegion)
+		require.Len(t, enumReps, 3)
+
+		expEnumReps := getEnumMembers(t, ctx, tc.Server(0), dbID)
+		require.Equal(t, len(expEnumReps), len(enumReps))
+
+		for r, rep := range expEnumReps {
+			res, ok := enumReps[catpb.RegionName(r)]
+			require.True(t, ok)
+			require.Equal(t, rep, res)
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}
+
 func getEnumMembers(
 	t *testing.T, ctx context.Context, ts serverutils.TestServerInterface, dbID descpb.ID,
 ) map[string][]byte {

pkg/server/server_sql.go

@@ -1306,7 +1306,7 @@ func (s *SQLServer) preStart(
 		}
 		// Start instance ID reclaim loop.
 		if err := s.sqlInstanceStorage.RunInstanceIDReclaimLoop(
-			ctx, stopper, timeutil.DefaultTimeSource{}, session.Expiration,
+			ctx, stopper, timeutil.DefaultTimeSource{}, s.internalExecutorFactory, session.Expiration,
 		); err != nil {
 			return err
 		}

pkg/sql/region_util.go

@@ -1514,7 +1514,7 @@ func GetLocalityRegionEnumPhysicalRepresentation(
 	) error {
 		enumReps, primaryRegion = nil, "" // reset for retry
 		var err error
-		enumReps, primaryRegion, err = getRegionEnumRepresentations(ctx, txn, dbID, descsCol)
+		enumReps, primaryRegion, err = GetRegionEnumRepresentations(ctx, txn, dbID, descsCol)
 		return err
 	}); err != nil {
 		return nil, err
@@ -1534,11 +1534,11 @@ func GetLocalityRegionEnumPhysicalRepresentation(
 	return nil, errors.AssertionFailedf("primary region not found")
 }
 
-// getRegionEnumRepresentations returns representations stored in the
+// GetRegionEnumRepresentations returns representations stored in the
 // multi-region enum type associated with dbID, and the primary region of it.
 // An ErrNotMultiRegionDatabase error will be returned if the database isn't
 // multi-region.
-func getRegionEnumRepresentations(
+func GetRegionEnumRepresentations(
 	ctx context.Context, txn *kv.Txn, dbID descpb.ID, descsCol *descs.Collection,
 ) (enumReps map[catpb.RegionName][]byte, primaryRegion catpb.RegionName, err error) {
 	dbDesc, regionEnumDesc, err := getDBAndRegionEnumDescs(

pkg/sql/sqlinstance/instancestorage/BUILD.bazel

@@ -21,8 +21,10 @@ go_library(
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",
+        "//pkg/sql",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/enum",
         "//pkg/sql/rowenc/valueside",

pkg/sql/sqlinstance/instancestorage/instancestorage.go

@@ -26,7 +26,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/enum"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slstorage"
@@ -235,9 +237,6 @@ func (s *Storage) CreateInstance(
 // getAvailableInstanceIDForRegion retrieves an available instance ID for the
 // current region associated with Storage s, and returns errNoPreallocatedRows
 // if there are no available rows.
-//
-// TODO(jaylim-crl): Store current region enum in s once we implement regional
-// by row for the sql_instances table.
 func (s *Storage) getAvailableInstanceIDForRegion(
 	ctx context.Context, region []byte, txn *kv.Txn,
 ) (base.SQLInstanceID, error) {
@@ -387,10 +386,35 @@ func (s *Storage) RunInstanceIDReclaimLoop(
 	ctx context.Context,
 	stopper *stop.Stopper,
 	ts timeutil.TimeSource,
+	internalExecutorFactory descs.TxnManager,
 	sessionExpirationFn func() hlc.Timestamp,
 ) error {
-	// TODO(jeffswenson): load regions from the system database enum.
-	regions := [][]byte{enum.One}
+	loadRegions := func() ([][]byte, error) {
+		// Load regions from the system DB.
+		var regions [][]byte
+		if err := internalExecutorFactory.DescsTxn(ctx, s.db, func(
+			ctx context.Context, txn *kv.Txn, descsCol *descs.Collection,
+		) error {
+			enumReps, _, err := sql.GetRegionEnumRepresentations(ctx, txn, keys.SystemDatabaseID, descsCol)
+			if err != nil {
+				if errors.Is(err, sql.ErrNotMultiRegionDatabase) {
+					return nil
+				}
+				return err
+			}
+			for _, r := range enumReps {
+				regions = append(regions, r)
+			}
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+		// The system database isn't multi-region.
+		if len(regions) == 0 {
+			regions = [][]byte{enum.One}
+		}
+		return regions, nil
+	}
 
 	return stopper.RunAsyncTask(ctx, "instance-id-reclaim-loop", func(ctx context.Context) {
 		ctx, cancel := stopper.WithCancelOnQuiesce(ctx)
@@ -411,6 +435,14 @@ func (s *Storage) RunInstanceIDReclaimLoop(
 			case <-timer.Ch():
 				timer.MarkRead()
 
+				// Load the regions each time we attempt to generate rows since
+				// regions can be added/removed to/from the system DB.
+				regions, err := loadRegions()
+				if err != nil {
+					log.Warningf(ctx, "failed to load regions from the system DB: %v", err)
+					continue
+				}
+
 				// Mark instances that belong to expired sessions as available
 				// and delete surplus IDs. Cleaning up surplus IDs is necessary
 				// to avoid ID exhaustion.
@@ -432,13 +464,6 @@ func (s *Storage) RunInstanceIDReclaimLoop(
 // generateAvailableInstanceRows allocates available instance IDs, and store
 // them in the sql_instances table. When instance IDs are pre-allocated, all
 // other fields in that row will be NULL.
-//
-// TODO(jaylim-crl): Handle multiple regions in this logic. encodeRow has to be
-// updated with crdb_region. When we handle multiple regions, we have to figure
-// out where to get the list of regions, and ensure that we don't do a global
-// read for each region assuming that the number of pre-allocated entries is
-// insufficient. One global KV read and write would be sufficient for **all**
-// regions.
 func (s *Storage) generateAvailableInstanceRows(
 	ctx context.Context, regions [][]byte, sessionExpiration hlc.Timestamp,
 ) error {