diff --git a/pkg/sql/sqlliveness/slinstance/BUILD.bazel b/pkg/sql/sqlliveness/slinstance/BUILD.bazel index 76d8b4216807..00f0f3674acf 100644 --- a/pkg/sql/sqlliveness/slinstance/BUILD.bazel +++ b/pkg/sql/sqlliveness/slinstance/BUILD.bazel @@ -43,6 +43,7 @@ go_test( "//pkg/sql/catalog", "//pkg/sql/catalog/descpb", "//pkg/sql/catalog/typedesc", + "//pkg/sql/enum", "//pkg/sql/sqlliveness", "//pkg/sql/sqlliveness/slstorage", "//pkg/util/hlc", diff --git a/pkg/sql/sqlliveness/slinstance/slinstance.go b/pkg/sql/sqlliveness/slinstance/slinstance.go index 987f48562cfa..970c21c0df98 100644 --- a/pkg/sql/sqlliveness/slinstance/slinstance.go +++ b/pkg/sql/sqlliveness/slinstance/slinstance.go @@ -200,8 +200,11 @@ func (l *Instance) clearSessionLocked(ctx context.Context) (createNewSession boo // createSession tries until it can create a new session and returns an error // only if the heart beat loop should exit. func (l *Instance) createSession(ctx context.Context) (*session, error) { - // TODO(jaylim-crl): Check l.currentRegion. Use enum.One if that is empty. - id, err := slstorage.MakeSessionID(enum.One, uuid.MakeV4()) + region := enum.One + if l.currentRegion != nil { + region = l.currentRegion + } + id, err := slstorage.MakeSessionID(region, uuid.MakeV4()) if err != nil { return nil, err } diff --git a/pkg/sql/sqlliveness/slinstance/slinstance_test.go b/pkg/sql/sqlliveness/slinstance/slinstance_test.go index d7941f7f3d04..73ff714b9c2a 100644 --- a/pkg/sql/sqlliveness/slinstance/slinstance_test.go +++ b/pkg/sql/sqlliveness/slinstance/slinstance_test.go @@ -16,7 +16,12 @@ import ( "time" "github.com/cockroachdb/cockroach/pkg/clusterversion" + "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/settings/cluster" + "github.com/cockroachdb/cockroach/pkg/sql/catalog" + "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" + "github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc" + "github.com/cockroachdb/cockroach/pkg/sql/enum" "github.com/cockroachdb/cockroach/pkg/sql/sqlliveness" "github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance" "github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slstorage" @@ -57,6 +62,11 @@ func TestSQLInstance(t *testing.T) { require.NoError(t, err) require.True(t, a) + region, id, err := slstorage.UnsafeDecodeSessionID(s1.ID()) + require.NoError(t, err) + require.Equal(t, enum.One, region) + require.NotNil(t, id) + s2, err := sqlInstance.Session(ctx) require.NoError(t, err) require.Equal(t, s1.ID(), s2.ID()) @@ -91,3 +101,59 @@ func TestSQLInstance(t *testing.T) { _, err = sqlInstance.Session(ctx) require.Error(t, err) } + +func TestSQLInstanceWithRegion(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx, stopper := context.Background(), stop.NewStopper() + defer stopper.Stop(ctx) + + clock := hlc.NewClock(timeutil.NewManualTime(timeutil.Unix(0, 42)), time.Nanosecond /* maxOffset */) + settings := cluster.MakeTestingClusterSettingsWithVersions( + clusterversion.TestingBinaryVersion, + clusterversion.TestingBinaryMinSupportedVersion, + true /* initializeVersion */) + slinstance.DefaultTTL.Override(ctx, &settings.SV, 20*time.Millisecond) + slinstance.DefaultHeartBeat.Override(ctx, &settings.SV, 10*time.Millisecond) + + mrDesc, _ := typedesc.NewBuilder((&descpb.TypeDescriptor{ + Kind: descpb.TypeDescriptor_MULTIREGION_ENUM, + RegionConfig: &descpb.TypeDescriptor_RegionConfig{ + PrimaryRegion: "us-east2", + }, + EnumMembers: []descpb.TypeDescriptor_EnumMember{ + { + LogicalRepresentation: "us-east1", + PhysicalRepresentation: []byte{1}, + Capability: descpb.TypeDescriptor_EnumMember_READ_ONLY, + }, + { + LogicalRepresentation: "us-east2", + PhysicalRepresentation: []byte{2}, + }, + { + LogicalRepresentation: "us-east3", + PhysicalRepresentation: []byte{3}, + }, + }, + })).BuildImmutable().(catalog.TypeDescriptor) + + fakeStorage := slstorage.NewFakeStorage() + sqlInstance := slinstance.NewSQLInstance(stopper, clock, fakeStorage, settings, nil, nil) + require.NoError(t, sqlInstance.SetRegionalData(roachpb.Locality{Tiers: []roachpb.Tier{ + {Key: "region", Value: "us-east3"}, + }}, mrDesc)) + sqlInstance.Start(ctx) + + s1, err := sqlInstance.Session(ctx) + require.NoError(t, err) + a, err := fakeStorage.IsAlive(ctx, s1.ID()) + require.NoError(t, err) + require.True(t, a) + + region, id, err := slstorage.UnsafeDecodeSessionID(s1.ID()) + require.NoError(t, err) + require.Equal(t, []byte{3}, region) + require.NotNil(t, id) +}