diff --git a/conf/config.toml b/conf/config.toml index 252d5a2938e..488275267d6 100644 --- a/conf/config.toml +++ b/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 = [] diff --git a/server/config/config.go b/server/config/config.go index 6afef769485..8d483117fd1 100644 --- a/server/config/config.go +++ b/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"` +} diff --git a/server/keyspace/keyspace.go b/server/keyspace/keyspace.go index e9c4f22e86b..107c42a7952 100644 --- a/server/keyspace/keyspace.go +++ b/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 } diff --git a/server/keyspace/keyspace_test.go b/server/keyspace/keyspace_test.go index 472fff851a5..401e573ea4b 100644 --- a/server/keyspace/keyspace_test.go +++ b/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()) } diff --git a/server/server.go b/server/server.go index 260e226996e..d658115f526 100644 --- a/server/server.go +++ b/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( diff --git a/tests/server/keyspace/keyspace_test.go b/tests/server/keyspace/keyspace_test.go index 69b17be08b5..5ade284c2be 100644 --- a/tests/server/keyspace/keyspace_test.go +++ b/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,6 +42,9 @@ 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)) } @@ -48,7 +52,9 @@ func TestKeyspaceTestSuite(t *testing.T) { 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) + } +}