Skip to content

Commit

Permalink
sqlliveness: embed current region into session ID
Browse files Browse the repository at this point in the history
Follow up on #90408 and #91019.

Previously, the sqlliveness subsystem was using enum.One when constructing the
session ID since it doesn't have the regional data yet. This commit implements
that missing part of it by ensuring that the regional representation is
plumbed all the way to the sqlliveness subsystem whenever the SQL pod is
started, and use it to construct the session ID. This will enable our REGIONAL
BY ROW work to put the data in the right region.

The multi-region enum for the system database will be read during startup for
that to work. Since doing this already requires loading the system DB's
descriptor (which indirectly tests whether the system DB has been bootstrapped)
for the tenant, we can remove the call to checkTenantExists.

Epic: None

Release note: None
  • Loading branch information
jaylim-crl committed Nov 23, 2022
1 parent 979dfb0 commit 3685865
Show file tree
Hide file tree
Showing 11 changed files with 396 additions and 87 deletions.
3 changes: 3 additions & 0 deletions pkg/ccl/multiregionccl/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ go_test(
"multiregion_system_table_test.go",
"multiregion_test.go",
"region_test.go",
"region_util_test.go",
"regional_by_row_system_database_test.go",
"regional_by_row_test.go",
"roundtrips_test.go",
Expand Down Expand Up @@ -67,11 +68,13 @@ go_test(
"//pkg/sql/catalog/desctestutils",
"//pkg/sql/catalog/tabledesc",
"//pkg/sql/catalog/typedesc",
"//pkg/sql/enum",
"//pkg/sql/execinfra",
"//pkg/sql/parser",
"//pkg/sql/rowenc",
"//pkg/sql/sem/tree",
"//pkg/sql/sqlliveness",
"//pkg/sql/sqlliveness/slstorage",
"//pkg/sql/sqltestutils",
"//pkg/sql/tests",
"//pkg/sql/types",
Expand Down
97 changes: 89 additions & 8 deletions pkg/ccl/multiregionccl/multiregion_system_table_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
"github.com/cockroachdb/cockroach/pkg/sql/enum"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slstorage"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
Expand Down Expand Up @@ -176,10 +178,10 @@ func TestMrSystemDatabase(t *testing.T) {

t.Run("InUse", func(t *testing.T) {
query := `
SELECT id, addr, session_id, locality, crdb_region
FROM system.sql_instances
WHERE session_id IS NOT NULL
`
SELECT id, addr, session_id, locality, crdb_region
FROM system.sql_instances
WHERE session_id IS NOT NULL
`
rows := tDB.Query(t, query)
require.True(t, rows.Next())
for {
Expand Down Expand Up @@ -207,10 +209,10 @@ func TestMrSystemDatabase(t *testing.T) {

t.Run("Preallocated", func(t *testing.T) {
query := `
SELECT id, addr, session_id, locality, crdb_region
FROM system.sql_instances
WHERE session_id IS NULL
`
SELECT id, addr, session_id, locality, crdb_region
FROM system.sql_instances
WHERE session_id IS NULL
`
rows := tDB.Query(t, query)
require.True(t, rows.Next())
for {
Expand All @@ -234,3 +236,82 @@ func TestMrSystemDatabase(t *testing.T) {
})
})
}

func TestTenantStartupWithMultiRegionEnum(t *testing.T) {
defer leaktest.AfterTest(t)()
defer log.Scope(t).Close(t)
defer envutil.TestSetEnv(t, "COCKROACH_MR_SYSTEM_DATABASE", "1")()

// Enable settings required for configuring a tenant's system database as multi-region.
cs := cluster.MakeTestingClusterSettings()
sql.SecondaryTenantsMultiRegionAbstractionsEnabled.Override(context.Background(), &cs.SV, true)
sql.SecondaryTenantZoneConfigsEnabled.Override(context.Background(), &cs.SV, true)

tc, _, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
t, 3 /*numServers*/, base.TestingKnobs{}, multiregionccltestutils.WithSettings(cs),
)
defer cleanup()

tenID := roachpb.MustMakeTenantID(10)
ten, tSQL := serverutils.StartTenant(t, tc.Server(0), base.TestTenantArgs{
Settings: cs,
TenantID: tenID,
Locality: roachpb.Locality{
Tiers: []roachpb.Tier{
{Key: "region", Value: "us-east1"},
},
},
})
defer tSQL.Close()
tenSQLDB := sqlutils.MakeSQLRunner(tSQL)

// Update system database with regions.
tenSQLDB.Exec(t, `ALTER DATABASE system SET PRIMARY REGION "us-east1"`)
tenSQLDB.Exec(t, `ALTER DATABASE system ADD REGION "us-east2"`)
tenSQLDB.Exec(t, `ALTER DATABASE system ADD REGION "us-east3"`)

ten2, tSQL2 := serverutils.StartTenant(t, tc.Server(2), base.TestTenantArgs{
Settings: cs,
TenantID: tenID,
Existing: true,
Locality: roachpb.Locality{
Tiers: []roachpb.Tier{
{Key: "region", Value: "us-east3"},
},
},
})
defer tSQL2.Close()
tenSQLDB2 := sqlutils.MakeSQLRunner(tSQL2)

// The sqlliveness entry created by the first SQL server has enum.One as the
// region as the system database hasn't been updated when it first started.
var sessionID string
tenSQLDB2.QueryRow(t, `SELECT session_id FROM system.sql_instances WHERE id = $1`,
ten.SQLInstanceID()).Scan(&sessionID)
region, id, err := slstorage.UnsafeDecodeSessionID(sqlliveness.SessionID(sessionID))
require.NoError(t, err)
require.NotNil(t, id)
require.Equal(t, enum.One, region)

// Ensure that the sqlliveness entry created by the second SQL server has
// the right region and session UUID.
tenSQLDB2.QueryRow(t, `SELECT session_id FROM system.sql_instances WHERE id = $1`,
ten2.SQLInstanceID()).Scan(&sessionID)
region, id, err = slstorage.UnsafeDecodeSessionID(sqlliveness.SessionID(sessionID))
require.NoError(t, err)
require.NotNil(t, id)
require.NotEqual(t, enum.One, region)

rows := tenSQLDB2.Query(t, `SELECT crdb_region, session_uuid FROM system.sqlliveness`)
defer rows.Close()
livenessMap := map[string][]byte{}
for rows.Next() {
var region, sessionUUID string
require.NoError(t, rows.Scan(&region, &sessionUUID))
livenessMap[sessionUUID] = []byte(region)
}
require.NoError(t, rows.Err())
r, ok := livenessMap[string(id)]
require.True(t, ok)
require.Equal(t, r, region)
}
108 changes: 108 additions & 0 deletions pkg/ccl/multiregionccl/region_util_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
// Copyright 2022 The Cockroach Authors.
//
// Licensed as a CockroachDB Enterprise file under the Cockroach Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt

package multiregionccl_test

import (
"context"
"testing"

"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/ccl/multiregionccl/multiregionccltestutils"
"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/descpb"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/stretchr/testify/require"
)

// TestGetLocalityRegionEnumPhysicalRepresentation is in the ccl package since
// it utilizes adding regions to a database.
func TestGetLocalityRegionEnumPhysicalRepresentation(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"`)

s0 := tc.ServerTyped(0)
ief := s0.InternalExecutorFactory().(descs.TxnManager)
dbID := 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{
Tiers: []roachpb.Tier{{Key: "region", Value: "us-east2"}},
},
)
require.NoError(t, err)

enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(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{
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))
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{})
require.NoError(t, err)

// Fallback to primary region is locality information is missing.
enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
require.NotEmpty(t, enumMembers)
require.Equal(t, enumMembers["us-east1"], regionEnum)
})
}

func getEnumMembers(
t *testing.T, ctx context.Context, ts serverutils.TestServerInterface, dbID descpb.ID,
) map[string][]byte {
t.Helper()
enumMembers := make(map[string][]byte)
err := sql.TestingDescsTxn(ctx, ts, func(ctx context.Context, txn *kv.Txn, descsCol *descs.Collection) error {
_, dbDesc, err := descsCol.GetImmutableDatabaseByID(ctx, txn, dbID,
tree.DatabaseLookupFlags{Required: true})
require.NoError(t, err)
regionEnumID, err := dbDesc.MultiRegionEnumID()
require.NoError(t, err)
regionEnumDesc, err := descsCol.GetImmutableTypeByID(ctx, txn, regionEnumID,
tree.ObjectLookupFlags{CommonLookupFlags: tree.CommonLookupFlags{Required: true}})
require.NoError(t, err)
for ord := 0; ord < regionEnumDesc.NumEnumMembers(); ord++ {
enumMembers[regionEnumDesc.GetMemberLogicalRepresentation(ord)] = regionEnumDesc.GetMemberPhysicalRepresentation(ord)
}
return nil
})
require.NoError(t, err)
return enumMembers
}
3 changes: 1 addition & 2 deletions pkg/ccl/serverccl/server_sql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,8 +198,7 @@ func TestNonExistentTenant(t *testing.T) {
DisableCreateTenant: true,
SkipTenantCheck: true,
})
require.Error(t, err)
require.Equal(t, "system DB uninitialized, check if tenant is non existent", err.Error())
require.EqualError(t, err, `database "[1]" does not exist`)
}

// TestTenantRowIDs confirms `unique_rowid()` works as expected in a
Expand Down
42 changes: 10 additions & 32 deletions pkg/server/server_sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,6 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/hydrateddesc"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
"github.com/cockroachdb/cockroach/pkg/sql/colexec"
"github.com/cockroachdb/cockroach/pkg/sql/consistencychecker"
"github.com/cockroachdb/cockroach/pkg/sql/contention"
Expand Down Expand Up @@ -1256,28 +1255,6 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
}, nil
}

// Checks if tenant exists. This function does a very superficial check to see if the system db
// has been bootstrapped for the tenant. This is not a complete check and is only sufficient
// to be used in the dev environment.
func checkTenantExists(ctx context.Context, codec keys.SQLCodec, db *kv.DB) error {
if codec.ForSystemTenant() {
return errors.AssertionFailedf("asked to check for tenant but system codec specified")
}

key := catalogkeys.MakeDatabaseNameKey(codec, systemschema.SystemDatabaseName)
result, err := db.Get(ctx, key)
if err != nil {
return err
}
if result.Value == nil || result.ValueInt() != keys.SystemDatabaseID {
return errors.New("system DB uninitialized, check if tenant is non existent")
}
// Tenant has been confirmed to be bootstrapped successfully
// as the system database, which is a part of the bootstrap data for
// a tenant keyspace, exists in the namespace table.
return nil
}

func (s *SQLServer) setInstanceID(
ctx context.Context, instanceID base.SQLInstanceID, sessionID sqlliveness.SessionID,
) error {
Expand All @@ -1297,7 +1274,6 @@ func (s *SQLServer) preStart(
pgL net.Listener,
orphanedLeasesTimeThresholdNanos int64,
) error {

// If necessary, start the tenant proxy first, to ensure all other
// components can properly route to KV nodes. The Start method will block
// until a connection is established to the cluster and its ID has been
Expand All @@ -1306,17 +1282,19 @@ func (s *SQLServer) preStart(
if err := s.tenantConnect.Start(ctx); err != nil {
return err
}
// Confirm tenant exists prior to initialization. This is a sanity
// check for the dev environment to ensure that a tenant has been
// successfully created before attempting to initialize a SQL
// server for it.
if err := checkTenantExists(ctx, s.execCfg.Codec, s.execCfg.DB); err != nil {
return err
}
}

// Load the multi-region enum by reading the system database's descriptor.
// This also serves as a simple check to see if a tenant exist (i.e. by
// checking whether the system db has been bootstrapped).
regionPhysicalRep, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
ctx, s.internalExecutorFactory, s.execCfg.DB, keys.SystemDatabaseID, s.distSQLServer.Locality)
if err != nil && !errors.Is(err, sql.ErrNotMultiRegionDatabase) {
return err
}

// Start the sql liveness subsystem. We'll need it to get a session.
s.sqlLivenessProvider.Start(ctx)
s.sqlLivenessProvider.Start(ctx, regionPhysicalRep)

_, isMixedSQLAndKVNode := s.sqlIDContainer.OptionalNodeID()
isTenant := !isMixedSQLAndKVNode
Expand Down
Loading

0 comments on commit 3685865

Please sign in to comment.