keyspace: Pre alloc keyspace (#5869)

ref #5293

Signed-off-by: David <8039876+AmoebaProtozoa@users.noreply.github.com>

conf/config.toml

@@ -200,3 +200,10 @@
 
 ## When enabled, usage data will be sent to PingCAP for improving user experience.
 # enable-telemetry = true
+
+[keyspaces]
+## pre-alloc is used to pre-allocate keyspaces during pd bootstrap.
+## Its value should be a list of strings, denotting the name of the keyspaces.
+## Example:
+## pre-alloc = ["admin", "user1", "user2"]
+# pre-alloc = []

server/config/config.go

@@ -165,6 +165,8 @@ type Config struct {
 	Dashboard DashboardConfig `toml:"dashboard" json:"dashboard"`
 
 	ReplicationMode ReplicationModeConfig `toml:"replication-mode" json:"replication-mode"`
+
+	Keyspace KeyspaceConfig `toml:"keyspace" json:"keyspace"`
 }
 
 // NewConfig creates a new config.
@@ -1477,3 +1479,9 @@ type SecurityConfig struct {
 	RedactInfoLog bool              `toml:"redact-info-log" json:"redact-info-log"`
 	Encryption    encryption.Config `toml:"encryption" json:"encryption"`
 }
+
+// KeyspaceConfig is the configuration for keyspace management.
+type KeyspaceConfig struct {
+	// PreAlloc contains the keyspace to be allocated during keyspace manager initialization.
+	PreAlloc []string `toml:"pre-alloc" json:"pre-alloc"`
+}

server/keyspace/keyspace.go

@@ -28,6 +28,7 @@ import (
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/tikv/pd/server/cluster"
+	"github.com/tikv/pd/server/config"
 	"go.uber.org/zap"
 )
 
@@ -60,6 +61,8 @@ type Manager struct {
 	rc *cluster.RaftCluster
 	// ctx is the context of the manager, to be used in transaction.
 	ctx context.Context
+	// config is the configurations of the manager.
+	config config.KeyspaceConfig
 }
 
 // CreateKeyspaceRequest represents necessary arguments to create a keyspace.
@@ -73,13 +76,18 @@ type CreateKeyspaceRequest struct {
 }
 
 // NewKeyspaceManager creates a Manager of keyspace related data.
-func NewKeyspaceManager(store endpoint.KeyspaceStorage, rc *cluster.RaftCluster, idAllocator id.Allocator) *Manager {
+func NewKeyspaceManager(store endpoint.KeyspaceStorage,
+	rc *cluster.RaftCluster,
+	idAllocator id.Allocator,
+	config config.KeyspaceConfig,
+) *Manager {
 	return &Manager{
 		metaLock:    syncutil.NewLockGroup(syncutil.WithHash(keyspaceIDHash)),
 		idAllocator: idAllocator,
 		store:       store,
 		rc:          rc,
 		ctx:         context.TODO(),
+		config:      config,
 	}
 }
 
@@ -103,6 +111,19 @@ func (manager *Manager) Bootstrap() error {
 	if err != nil && err != ErrKeyspaceExists {
 		return err
 	}
+
+	// Initialize pre-alloc keyspace.
+	preAlloc := manager.config.PreAlloc
+	for _, keyspaceName := range preAlloc {
+		_, err = manager.CreateKeyspace(&CreateKeyspaceRequest{
+			Name: keyspaceName,
+			Now:  now,
+		})
+		// Ignore the keyspaceExists error for the same reason as saving default keyspace.
+		if err != nil && err != ErrKeyspaceExists {
+			return err
+		}
+	}
 	return nil
 }
 

server/keyspace/keyspace_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/tikv/pd/pkg/mock/mockid"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
+	"github.com/tikv/pd/server/config"
 )
 
 const (
@@ -49,7 +50,7 @@ func TestKeyspaceTestSuite(t *testing.T) {
 func (suite *keyspaceTestSuite) SetupTest() {
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
 	allocator := mockid.NewIDAllocator()
-	suite.manager = NewKeyspaceManager(store, nil, allocator)
+	suite.manager = NewKeyspaceManager(store, nil, allocator, config.KeyspaceConfig{})
 	suite.NoError(suite.manager.Bootstrap())
 }
 

server/server.go

@@ -406,7 +406,7 @@ func (s *Server) startServer(ctx context.Context) error {
 		Member:    s.member.MemberValue(),
 		Step:      keyspace.AllocStep,
 	})
-	s.keyspaceManager = keyspace.NewKeyspaceManager(s.storage, s.cluster, keyspaceIDAllocator)
+	s.keyspaceManager = keyspace.NewKeyspaceManager(s.storage, s.cluster, keyspaceIDAllocator, s.cfg.Keyspace)
 	s.hbStreams = hbstream.NewHeartbeatStreams(ctx, s.clusterID, s.cluster)
 	// initial hot_region_storage in here.
 	s.hotRegionStorage, err = storage.NewHotRegionsStorage(

tests/server/keyspace/keyspace_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/pd/pkg/codec"
 	"github.com/tikv/pd/pkg/storage/endpoint"
+	"github.com/tikv/pd/server/config"
 	"github.com/tikv/pd/server/keyspace"
 	"github.com/tikv/pd/server/schedule/labeler"
 	"github.com/tikv/pd/tests"
@@ -41,14 +42,19 @@ type keyspaceTestSuite struct {
 	manager *keyspace.Manager
 }
 
+// preAllocKeyspace is used to test keyspace pre-allocation.
+var preAllocKeyspace = []string{"pre-alloc0", "pre-alloc1", "pre-alloc2"}
+
 func TestKeyspaceTestSuite(t *testing.T) {
 	suite.Run(t, new(keyspaceTestSuite))
 }
 
 func (suite *keyspaceTestSuite) SetupTest() {
 	ctx, cancel := context.WithCancel(context.Background())
 	suite.cleanup = cancel
-	cluster, err := tests.NewTestCluster(ctx, 3)
+	cluster, err := tests.NewTestCluster(ctx, 3, func(conf *config.Config, serverName string) {
+		conf.Keyspace.PreAlloc = preAllocKeyspace
+	})
 	suite.cluster = cluster
 	suite.NoError(err)
 	suite.NoError(cluster.RunInitialServers())
@@ -109,3 +115,15 @@ func checkLabelRule(re *require.Assertions, id uint32, regionLabeler *labeler.Re
 	re.Equal(txnLeftBound, rangeRule[1].StartKeyHex)
 	re.Equal(txnRightBound, rangeRule[1].EndKeyHex)
 }
+
+func (suite *keyspaceTestSuite) TestPreAlloc() {
+	re := suite.Require()
+	regionLabeler := suite.server.GetRaftCluster().GetRegionLabeler()
+	for _, keyspaceName := range preAllocKeyspace {
+		// Check pre-allocated keyspaces are correctly allocated.
+		meta, err := suite.manager.LoadKeyspace(keyspaceName)
+		re.NoError(err)
+		// Check pre-allocated keyspaces also have the correct region label.
+		checkLabelRule(re, meta.GetId(), regionLabeler)
+	}
+}