*: define user kind for keyspace group

pkg/keyspace/keyspace.go

@@ -46,6 +46,10 @@ const (
 	regionLabelIDPrefix = "keyspaces/"
 	// regionLabelKey is the key for keyspace id in keyspace region label.
 	regionLabelKey = "id"
+	// UserKindKey is the key for user kind in keyspace config.
+	UserKindKey = "user_kind"
+	// TSOKeyspaceGroupIDKey is the key for tso keyspace group id in keyspace config.
+	TSOKeyspaceGroupIDKey = "tso_keyspace_group_id"
 )
 
 // Config is the interface for keyspace config.
@@ -68,6 +72,7 @@ type Manager struct {
 	ctx context.Context
 	// config is the configurations of the manager.
 	config Config
+	kgm    *GroupManager
 }
 
 // CreateKeyspaceRequest represents necessary arguments to create a keyspace.
@@ -85,6 +90,7 @@ func NewKeyspaceManager(store endpoint.KeyspaceStorage,
 	cluster schedule.Cluster,
 	idAllocator id.Allocator,
 	config Config,
+	kgm *GroupManager,
 ) *Manager {
 	return &Manager{
 		metaLock:    syncutil.NewLockGroup(syncutil.WithHash(keyspaceIDHash)),
@@ -93,6 +99,7 @@ func NewKeyspaceManager(store endpoint.KeyspaceStorage,
 		cluster:     cluster,
 		ctx:         context.TODO(),
 		config:      config,
+		kgm:         kgm,
 	}
 }
 
@@ -103,14 +110,22 @@ func (manager *Manager) Bootstrap() error {
 		return err
 	}
 	now := time.Now().Unix()
+	id, err := manager.kgm.GetAvailableKeyspaceGroupIDByKind(endpoint.Basic)
+	if err != nil {
+		return err
+	}
 	defaultKeyspace := &keyspacepb.KeyspaceMeta{
 		Id:             DefaultKeyspaceID,
 		Name:           DefaultKeyspaceName,
 		State:          keyspacepb.KeyspaceState_ENABLED,
 		CreatedAt:      now,
 		StateChangedAt: now,
+		Config: map[string]string{
+			UserKindKey:           endpoint.Basic.String(),
+			TSOKeyspaceGroupIDKey: id,
+		},
 	}
-	err := manager.saveNewKeyspace(defaultKeyspace)
+	err = manager.saveNewKeyspace(defaultKeyspace)
 	// It's possible that default keyspace already exists in the storage (e.g. PD restart/recover),
 	// so we ignore the keyspaceExists error.
 	if err != nil && err != ErrKeyspaceExists {
@@ -120,9 +135,17 @@ func (manager *Manager) Bootstrap() error {
 	// Initialize pre-alloc keyspace.
 	preAlloc := manager.config.GetPreAlloc()
 	for _, keyspaceName := range preAlloc {
+		id, err := manager.kgm.GetAvailableKeyspaceGroupIDByKind(endpoint.Basic)
+		if err != nil {
+			return err
+		}
 		_, err = manager.CreateKeyspace(&CreateKeyspaceRequest{
 			Name: keyspaceName,
 			Now:  now,
+			Config: map[string]string{
+				UserKindKey:           endpoint.Basic.String(),
+				TSOKeyspaceGroupIDKey: id,
+			},
 		})
 		// Ignore the keyspaceExists error for the same reason as saving default keyspace.
 		if err != nil && err != ErrKeyspaceExists {
@@ -148,6 +171,16 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac
 	if err != nil {
 		return nil, err
 	}
+	userKind := endpoint.StringUserKind(request.Config[UserKindKey])
+	id, err := manager.kgm.GetAvailableKeyspaceGroupIDByKind(userKind)
+	if err != nil {
+		return nil, err
+	}
+	if request.Config == nil {
+		request.Config = make(map[string]string)
+	}
+	request.Config[TSOKeyspaceGroupIDKey] = id
+	request.Config[UserKindKey] = userKind.String()
 	// Create and save keyspace metadata.
 	keyspace := &keyspacepb.KeyspaceMeta{
 		Id:             newID,

pkg/keyspace/keyspace_test.go

@@ -15,6 +15,7 @@
 package keyspace
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"strconv"
@@ -39,6 +40,8 @@ const (
 
 type keyspaceTestSuite struct {
 	suite.Suite
+	ctx     context.Context
+	cancel  context.CancelFunc
 	manager *Manager
 }
 
@@ -53,12 +56,18 @@ type mockConfig struct {
 func (m *mockConfig) GetPreAlloc() []string { return m.PreAlloc }
 
 func (suite *keyspaceTestSuite) SetupTest() {
+	suite.ctx, suite.cancel = context.WithCancel(context.Background())
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
 	allocator := mockid.NewIDAllocator()
-	suite.manager = NewKeyspaceManager(store, nil, allocator, &mockConfig{})
+	kgm := NewKeyspaceGroupManager(suite.ctx, store)
+	suite.manager = NewKeyspaceManager(store, nil, allocator, &mockConfig{}, kgm)
 	suite.NoError(suite.manager.Bootstrap())
 }
 
+func (suite *keyspaceTestSuite) TearDownTest() {
+	suite.cancel()
+}
+
 func (suite *keyspaceTestSuite) SetupSuite() {
 	suite.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)"))
 }
@@ -155,6 +164,9 @@ func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfig() {
 	// Changing config of DEFAULT keyspace is allowed.
 	updated, err := manager.UpdateKeyspaceConfig(DefaultKeyspaceName, mutations)
 	re.NoError(err)
+	// remove auto filled fields
+	delete(updated.Config, TSOKeyspaceGroupIDKey)
+	delete(updated.Config, UserKindKey)
 	checkMutations(re, nil, updated.Config, mutations)
 }
 

pkg/keyspace/tso_keyspace_group.go

@@ -40,9 +40,8 @@ func NewKeyspaceGroupManager(ctx context.Context, store endpoint.KeyspaceGroupSt
 // Bootstrap saves default keyspace group info.
 func (m *GroupManager) Bootstrap() error {
 	defaultKeyspaceGroup := &endpoint.KeyspaceGroup{
-		ID: utils.DefaultKeySpaceGroupID,
-		// TODO: define a user kind type
-		UserKind: "default",
+		ID:       utils.DefaultKeySpaceGroupID,
+		UserKind: endpoint.Basic.String(),
 	}
 	err := m.saveKeyspaceGroups([]*endpoint.KeyspaceGroup{defaultKeyspaceGroup})
 	// It's possible that default keyspace group already exists in the storage (e.g. PD restart/recover),
@@ -111,3 +110,9 @@ func (m *GroupManager) saveKeyspaceGroups(keyspaceGroups []*endpoint.KeyspaceGro
 		return nil
 	})
 }
+
+// GetAvailableKeyspaceGroupIDByKind returns the available keyspace group id by user kind.
+func (m *GroupManager) GetAvailableKeyspaceGroupIDByKind(userKind endpoint.UserKind) (string, error) {
+	// TODO: implement it
+	return "0", nil
+}

pkg/keyspace/tso_keyspace_group_test.go

@@ -25,6 +25,8 @@ import (
 
 type keyspaceGroupTestSuite struct {
 	suite.Suite
+	ctx     context.Context
+	cancel  context.CancelFunc
 	manager *GroupManager
 }
 
@@ -33,26 +35,31 @@ func TestKeyspaceGroupTestSuite(t *testing.T) {
 }
 
 func (suite *keyspaceGroupTestSuite) SetupTest() {
+	suite.ctx, suite.cancel = context.WithCancel(context.Background())
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	suite.manager = NewKeyspaceGroupManager(context.Background(), store)
+	suite.manager = NewKeyspaceGroupManager(suite.ctx, store)
 	suite.NoError(suite.manager.Bootstrap())
 }
 
+func (suite *keyspaceGroupTestSuite) TearDownTest() {
+	suite.cancel()
+}
+
 func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() {
 	re := suite.Require()
 
 	keyspaceGroups := []*endpoint.KeyspaceGroup{
 		{
 			ID:       uint32(1),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 		{
 			ID:       uint32(2),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 		{
 			ID:       uint32(3),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 	}
 	err := suite.manager.CreateKeyspaceGroups(keyspaceGroups)
@@ -69,11 +76,11 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() {
 	kg, err := suite.manager.GetKeyspaceGroupByID(0)
 	re.NoError(err)
 	re.Equal(uint32(0), kg.ID)
-	re.Equal("default", kg.UserKind)
+	re.Equal(endpoint.Basic.String(), kg.UserKind)
 	kg, err = suite.manager.GetKeyspaceGroupByID(3)
 	re.NoError(err)
 	re.Equal(uint32(3), kg.ID)
-	re.Equal("business", kg.UserKind)
+	re.Equal(endpoint.Standard.String(), kg.UserKind)
 	// remove the keyspace group 3
 	err = suite.manager.DeleteKeyspaceGroupByID(3)
 	re.NoError(err)
@@ -83,7 +90,7 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() {
 	re.Empty(kg)
 
 	// create an existing keyspace group
-	keyspaceGroups = []*endpoint.KeyspaceGroup{{ID: uint32(1), UserKind: "business"}}
+	keyspaceGroups = []*endpoint.KeyspaceGroup{{ID: uint32(1), UserKind: endpoint.Standard.String()}}
 	err = suite.manager.CreateKeyspaceGroups(keyspaceGroups)
 	re.Error(err)
 }

pkg/storage/endpoint/tso_keyspace_group.go

@@ -22,6 +22,44 @@ import (
 	"go.etcd.io/etcd/clientv3"
 )
 
+// UserKind represents the user kind.
+type UserKind int
+
+// Different user kinds.
+const (
+	Basic UserKind = iota
+	Standard
+	Enterprise
+
+	UserKindCount
+)
+
+// StringUserKind creates a UserKind with string.
+func StringUserKind(input string) UserKind {
+	switch input {
+	case Basic.String():
+		return Basic
+	case Standard.String():
+		return Standard
+	case Enterprise.String():
+		return Enterprise
+	default:
+		return Basic
+	}
+}
+
+func (k UserKind) String() string {
+	switch k {
+	case Basic:
+		return "basic"
+	case Standard:
+		return "standard"
+	case Enterprise:
+		return "enterprise"
+	}
+	return "unknown UserKind"
+}
+
 // KeyspaceGroupMember defines an election member which campaigns for the primary of the keyspace group.
 type KeyspaceGroupMember struct {
 	Address string `json:"address"`

server/server.go

@@ -435,8 +435,8 @@ 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.cfg.Keyspace)
 	s.keyspaceGroupManager = keyspace.NewKeyspaceGroupManager(s.ctx, s.storage)
+	s.keyspaceManager = keyspace.NewKeyspaceManager(s.storage, s.cluster, keyspaceIDAllocator, &s.cfg.Keyspace, s.keyspaceGroupManager)
 	s.hbStreams = hbstream.NewHeartbeatStreams(ctx, s.clusterID, s.cluster)
 	// initial hot_region_storage in here.
 	s.hotRegionStorage, err = storage.NewHotRegionsStorage(

tests/server/apiv2/handlers/keyspace_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/pd/pkg/keyspace"
+	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/server/apiv2/handlers"
 	"github.com/tikv/pd/tests"
@@ -218,6 +219,9 @@ func mustCreateKeyspace(re *require.Assertions, server *tests.TestServer, reques
 	re.NoError(err)
 	meta := &handlers.KeyspaceMeta{}
 	re.NoError(json.Unmarshal(data, meta))
+	// When creating a keyspace, it will be assigned a keyspace group id.
+	request.Config[keyspace.TSOKeyspaceGroupIDKey] = "0"
+	request.Config[keyspace.UserKindKey] = endpoint.Basic.String()
 	checkCreateRequest(re, request, meta.KeyspaceMeta)
 	return meta.KeyspaceMeta
 }

tests/server/apiv2/handlers/tso_keyspace_group_test.go

@@ -33,7 +33,8 @@ const keyspaceGroupsPrefix = "/pd/api/v2/tso/keyspace-groups"
 
 type keyspaceGroupTestSuite struct {
 	suite.Suite
-	cleanup func()
+	ctx     context.Context
+	cancel  context.CancelFunc
 	cluster *tests.TestCluster
 	server  *tests.TestServer
 }
@@ -43,9 +44,8 @@ func TestKeyspaceGroupTestSuite(t *testing.T) {
 }
 
 func (suite *keyspaceGroupTestSuite) SetupTest() {
-	ctx, cancel := context.WithCancel(context.Background())
-	suite.cleanup = cancel
-	cluster, err := tests.NewTestCluster(ctx, 1)
+	suite.ctx, suite.cancel = context.WithCancel(context.Background())
+	cluster, err := tests.NewTestCluster(suite.ctx, 1)
 	suite.cluster = cluster
 	suite.NoError(err)
 	suite.NoError(cluster.RunInitialServers())
@@ -55,7 +55,7 @@ func (suite *keyspaceGroupTestSuite) SetupTest() {
 }
 
 func (suite *keyspaceGroupTestSuite) TearDownTest() {
-	suite.cleanup()
+	suite.cancel()
 	suite.cluster.Destroy()
 }
 
@@ -64,11 +64,11 @@ func (suite *keyspaceGroupTestSuite) TestCreateKeyspaceGroups() {
 	kgs := &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{
 		{
 			ID:       uint32(1),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 		{
 			ID:       uint32(2),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 	}}
 
@@ -80,11 +80,11 @@ func (suite *keyspaceGroupTestSuite) TestLoadKeyspaceGroup() {
 	kgs := &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{
 		{
 			ID:       uint32(1),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 		{
 			ID:       uint32(2),
-			UserKind: "business",
+			UserKind: endpoint.Standard.String(),
 		},
 	}}