define user kind for keyspace group

Signed-off-by: Ryan Leung <[email protected]>

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"
+}
+
 // KeyspaceGroup is the keyspace group.
 type KeyspaceGroup struct {
 	ID       uint32 `json:"id"`

server/keyspace/keyspace.go

@@ -63,6 +63,8 @@ type Manager struct {
 	ctx context.Context
 	// config is the configurations of the manager.
 	config config.KeyspaceConfig
+	// keyspaceGroupManager manages keyspace group related data.
+	kgm *GroupManager
 }
 
 // CreateKeyspaceRequest represents necessary arguments to create a keyspace.
@@ -80,6 +82,7 @@ func NewKeyspaceManager(store endpoint.KeyspaceStorage,
 	rc *cluster.RaftCluster,
 	idAllocator id.Allocator,
 	config config.KeyspaceConfig,
+	kgm *GroupManager,
 ) *Manager {
 	return &Manager{
 		metaLock:    syncutil.NewLockGroup(syncutil.WithHash(keyspaceIDHash)),
@@ -88,6 +91,7 @@ func NewKeyspaceManager(store endpoint.KeyspaceStorage,
 		rc:          rc,
 		ctx:         context.TODO(),
 		config:      config,
+		kgm:         kgm,
 	}
 }
 
@@ -98,12 +102,17 @@ func (manager *Manager) Bootstrap() error {
 		return err
 	}
 	now := time.Now().Unix()
+	id := manager.kgm.GetAvailableKeyspaceGroupIDByKind(endpoint.Basic)
 	defaultKeyspace := &keyspacepb.KeyspaceMeta{
 		Id:             DefaultKeyspaceID,
 		Name:           DefaultKeyspaceName,
 		State:          keyspacepb.KeyspaceState_ENABLED,
 		CreatedAt:      now,
 		StateChangedAt: now,
+		Config: map[string]string{
+			"user_kind":             endpoint.Basic.String(),
+			"tso_keyspace_group_id": id,
+		},
 	}
 	err := manager.saveNewKeyspace(defaultKeyspace)
 	// It's possible that default keyspace already exists in the storage (e.g. PD restart/recover),
@@ -115,9 +124,14 @@ func (manager *Manager) Bootstrap() error {
 	// Initialize pre-alloc keyspace.
 	preAlloc := manager.config.PreAlloc
 	for _, keyspaceName := range preAlloc {
+		id := manager.kgm.GetAvailableKeyspaceGroupIDByKind(endpoint.Basic)
 		_, err = manager.CreateKeyspace(&CreateKeyspaceRequest{
 			Name: keyspaceName,
 			Now:  now,
+			Config: map[string]string{
+				"user_kind":             endpoint.Basic.String(),
+				"tso_keyspace_group_id": id,
+			},
 		})
 		// Ignore the keyspaceExists error for the same reason as saving default keyspace.
 		if err != nil && err != ErrKeyspaceExists {
@@ -143,6 +157,16 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac
 	if err != nil {
 		return nil, err
 	}
+	userKind := endpoint.StringUserKind(request.Config["user_kind"])
+	id, err := manager.kgm.GetAvailableKeyspaceGroupIDByKind(userKind)
+	if err != nil {
+		return nil, err
+	}
+	if request.Config == nil {
+		request.Config = make(map[string]string)
+	}
+	request.Config["tso_keyspace_group_id"] = id
+	request.Config["user_kind"] = userKind.String()
 	// Create and save keyspace metadata.
 	keyspace := &keyspacepb.KeyspaceMeta{
 		Id:             newID,

server/keyspace/keyspace_test.go

@@ -15,6 +15,7 @@
 package keyspace
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"strconv"
@@ -40,6 +41,8 @@ const (
 
 type keyspaceTestSuite struct {
 	suite.Suite
+	ctx     context.Context
+	cancel  context.CancelFunc
 	manager *Manager
 }
 
@@ -48,12 +51,18 @@ func TestKeyspaceTestSuite(t *testing.T) {
 }
 
 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, config.KeyspaceConfig{})
+	kgm := NewKeyspaceGroupManager(suite.ctx, store)
+	suite.manager = NewKeyspaceManager(store, nil, allocator, config.KeyspaceConfig{}, kgm)
 	suite.NoError(suite.manager.Bootstrap())
 }
 
+func (suite *keyspaceTestSuite) TearDownTest() {
+	suite.cancel()
+}
+
 func (suite *keyspaceTestSuite) SetupSuite() {
 	suite.NoError(failpoint.Enable("github.com/tikv/pd/server/keyspace/skipSplitRegion", "return(true)"))
 }
@@ -149,6 +158,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, "tso_keyspace_group_id")
+	delete(updated.Config, "user_kind")
 	checkMutations(re, nil, updated.Config, mutations)
 }
 

server/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
+}

server/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)
 }

server/server.go

@@ -440,8 +440,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(
@@ -698,14 +698,14 @@ func (s *Server) bootstrapCluster(req *pdpb.BootstrapRequest) (*pdpb.BootstrapRe
 		return nil, err
 	}
 
-	if err = s.GetKeyspaceManager().Bootstrap(); err != nil {
-		log.Warn("bootstrap keyspace manager failed", errs.ZapError(err))
-	}
-
 	if err = s.GetKeyspaceGroupManager().Bootstrap(); err != nil {
 		log.Warn("bootstrap keyspace group manager failed", errs.ZapError(err))
 	}
 
+	if err = s.GetKeyspaceManager().Bootstrap(); err != nil {
+		log.Warn("bootstrap keyspace manager failed", errs.ZapError(err))
+	}
+
 	return &pdpb.BootstrapResponse{
 		ReplicationStatus: s.cluster.GetReplicationMode().GetReplicationStatus(),
 	}, nil

tests/server/apiv2/handlers/keyspace_test.go

@@ -218,6 +218,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["tso_keyspace_group_id"] = "0"
+	request.Config["user_kind"] = "basic"
 	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(),
 		},
 	}}