Merge branch 'master' into extract-conn

pkg/keyspace/keyspace.go

@@ -49,10 +49,10 @@ const (
  UserKindKey = "user_kind"
  // TSOKeyspaceGroupIDKey is the key for tso keyspace group id in keyspace config.
  TSOKeyspaceGroupIDKey = "tso_keyspace_group_id"
- // keyspacePatrolBatchSize is the batch size for keyspace assignment patrol.
- // the limit of etcd txn op is 128, keyspacePatrolBatchSize need to be less than it.
+ // maxEtcdTxnOps is the max value of operations in an etcd txn. The default limit of etcd txn op is 128.
+ // We use 120 here to leave some space for other operations.
  // See: https://github.com/etcd-io/etcd/blob/d3e43d4de6f6d9575b489dd7850a85e37e0f6b6c/server/embed/config.go#L61
- keyspacePatrolBatchSize = 120
+ maxEtcdTxnOps = 120
 )
 
 // Config is the interface for keyspace config.
@@ -669,7 +669,7 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
  zap.Duration("cost", time.Since(start)),
  zap.Uint64("patrolled-keyspace-count", patrolledKeyspaceCount),
  zap.Uint64("assigned-keyspace-count", assignedKeyspaceCount),
- zap.Int("batch-size", keyspacePatrolBatchSize),
+ zap.Int("batch-size", maxEtcdTxnOps),
  zap.Uint32("current-start-id", currentStartID),
  zap.Uint32("next-start-id", nextStartID),
  )
@@ -688,7 +688,10 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
  if defaultKeyspaceGroup.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
- keyspaces, err := manager.store.LoadRangeKeyspace(txn, manager.nextPatrolStartID, keyspacePatrolBatchSize)
+ if defaultKeyspaceGroup.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
+ keyspaces, err := manager.store.LoadRangeKeyspace(txn, manager.nextPatrolStartID, maxEtcdTxnOps)
  if err != nil {
  return err
  }
@@ -698,9 +701,9 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
  currentStartID = keyspaces[0].GetId()
  nextStartID = keyspaces[keyspaceNum-1].GetId() + 1
  }
- // If there are less than `keyspacePatrolBatchSize` keyspaces,
+ // If there are less than `maxEtcdTxnOps` keyspaces,
  // we have reached the end of the keyspace list.
- moreToPatrol = keyspaceNum == keyspacePatrolBatchSize
+ moreToPatrol = keyspaceNum == maxEtcdTxnOps
  var (
  assigned = false
  keyspaceIDsToUnlock = make([]uint32, 0, keyspaceNum)
@@ -735,7 +738,7 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
  err = manager.store.SaveKeyspaceMeta(txn, ks)
  if err != nil {
  log.Error("[keyspace] failed to save keyspace meta during patrol",
- zap.Int("batch-size", keyspacePatrolBatchSize),
+ zap.Int("batch-size", maxEtcdTxnOps),
  zap.Uint32("current-start-id", currentStartID),
  zap.Uint32("next-start-id", nextStartID),
  zap.Uint32("keyspace-id", ks.Id), zap.Error(err))
@@ -747,7 +750,7 @@ func (manager *Manager) PatrolKeyspaceAssignment() error {
  err = manager.kgm.store.SaveKeyspaceGroup(txn, defaultKeyspaceGroup)
  if err != nil {
  log.Error("[keyspace] failed to save default keyspace group meta during patrol",
- zap.Int("batch-size", keyspacePatrolBatchSize),
+ zap.Int("batch-size", maxEtcdTxnOps),
  zap.Uint32("current-start-id", currentStartID),
  zap.Uint32("next-start-id", nextStartID), zap.Error(err))
  return err

pkg/keyspace/keyspace_test.go

@@ -405,7 +405,7 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignment() {
 func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() {
  re := suite.Require()
  // Create some keyspaces without any keyspace group.
- for i := 1; i < keyspacePatrolBatchSize*2+1; i++ {
+ for i := 1; i < maxEtcdTxnOps*2+1; i++ {
  now := time.Now().Unix()
  err := suite.manager.saveNewKeyspace(&keyspacepb.KeyspaceMeta{
  Id: uint32(i),
@@ -420,7 +420,7 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() {
  defaultKeyspaceGroup, err := suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID)
  re.NoError(err)
  re.NotNil(defaultKeyspaceGroup)
- for i := 1; i < keyspacePatrolBatchSize*2+1; i++ {
+ for i := 1; i < maxEtcdTxnOps*2+1; i++ {
  re.NotContains(defaultKeyspaceGroup.Keyspaces, uint32(i))
  }
  // Patrol the keyspace assignment.
@@ -430,7 +430,7 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() {
  defaultKeyspaceGroup, err = suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID)
  re.NoError(err)
  re.NotNil(defaultKeyspaceGroup)
- for i := 1; i < keyspacePatrolBatchSize*2+1; i++ {
+ for i := 1; i < maxEtcdTxnOps*2+1; i++ {
  re.Contains(defaultKeyspaceGroup.Keyspaces, uint32(i))
  }
 }

pkg/keyspace/tso_keyspace_group.go

@@ -17,6 +17,7 @@ package keyspace
 import (
  "context"
  "encoding/json"
+ "sort"
  "strconv"
  "strings"
  "sync"
@@ -340,6 +341,9 @@ func (m *GroupManager) saveKeyspaceGroups(keyspaceGroups []*endpoint.KeyspaceGro
  if oldKG.IsSplitting() && overwrite {
  return ErrKeyspaceGroupInSplit
  }
+ if oldKG.IsMerging() && overwrite {
+ return ErrKeyspaceGroupInMerging
+ }
  newKG := &endpoint.KeyspaceGroup{
  ID: keyspaceGroup.ID,
  UserKind: keyspaceGroup.UserKind,
@@ -415,6 +419,9 @@ func (m *GroupManager) updateKeyspaceForGroupLocked(userKind endpoint.UserKind,
  if kg.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
+ if kg.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
 
  changed := false
 
@@ -469,6 +476,9 @@ func (m *GroupManager) UpdateKeyspaceGroup(oldGroupID, newGroupID string, oldUse
  if oldKG.IsSplitting() || newKG.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
+ if oldKG.IsMerging() || newKG.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
 
  var updateOld, updateNew bool
  if !slice.Contains(newKG.Keyspaces, keyspaceID) {
@@ -516,6 +526,10 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3
  if splitSourceKg.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
+ // A keyspace group can not be split when it is in merging.
+ if splitSourceKg.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
  // Check if the source keyspace group has enough replicas.
  if len(splitSourceKg.Members) < utils.KeyspaceGroupDefaultReplicaCount {
  return ErrKeyspaceGroupNotEnoughReplicas
@@ -618,11 +632,7 @@ func (m *GroupManager) FinishSplitKeyspaceByID(splitTargetID uint32) error {
  if err != nil {
  return err
  }
- err = m.store.SaveKeyspaceGroup(txn, splitSourceKg)
- if err != nil {
- return err
- }
- return nil
+ return m.store.SaveKeyspaceGroup(txn, splitSourceKg)
  }); err != nil {
  return err
  }
@@ -663,6 +673,9 @@ func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, desiredReplicaCount
  if kg.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
+ if kg.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
  exists := make(map[string]struct{})
  for _, member := range kg.Members {
  exists[member.Address] = struct{}{}
@@ -719,6 +732,9 @@ func (m *GroupManager) SetNodesForKeyspaceGroup(id uint32, nodes []string) error
  if kg.IsSplitting() {
  return ErrKeyspaceGroupInSplit
  }
+ if kg.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
  members := make([]endpoint.KeyspaceGroupMember, 0, len(nodes))
  for _, node := range nodes {
  members = append(members, endpoint.KeyspaceGroupMember{Address: node})
@@ -743,3 +759,123 @@ func (m *GroupManager) IsExistNode(addr string) bool {
  }
  return false
 }
+
+// MergeKeyspaceGroups merges the keyspace group in the list into the target keyspace group.
+func (m *GroupManager) MergeKeyspaceGroups(mergeTargetID uint32, mergeList []uint32) error {
+ mergeListNum := len(mergeList)
+ if mergeListNum == 0 {
+ return nil
+ }
+ // The transaction below will:
+ // - Load and delete the keyspace groups in the merge list.
+ // - Load and update the target keyspace group.
+ // So we pre-check the number of operations to avoid exceeding the maximum number of etcd transaction.
+ if (mergeListNum+1)*2 > maxEtcdTxnOps {
+ return ErrExceedMaxEtcdTxnOps
+ }
+ if slice.Contains(mergeList, utils.DefaultKeyspaceGroupID) {
+ return ErrModifyDefaultKeyspaceGroup
+ }
+ var (
+ groups = make(map[uint32]*endpoint.KeyspaceGroup, mergeListNum+1)
+ mergeTargetKg *endpoint.KeyspaceGroup
+ )
+ m.Lock()
+ defer m.Unlock()
+ if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) (err error) {
+ // Load and check all keyspace groups first.
+ for _, kgID := range append(mergeList, mergeTargetID) {
+ kg, err := m.store.LoadKeyspaceGroup(txn, kgID)
+ if err != nil {
+ return err
+ }
+ if kg == nil {
+ return ErrKeyspaceGroupNotExists
+ }
+ // A keyspace group can not be merged if it's in splitting.
+ if kg.IsSplitting() {
+ return ErrKeyspaceGroupInSplit
+ }
+ // A keyspace group can not be split when it is in merging.
+ if kg.IsMerging() {
+ return ErrKeyspaceGroupInMerging
+ }
+ groups[kgID] = kg
+ }
+ mergeTargetKg = groups[mergeTargetID]
+ keyspaces := make(map[uint32]struct{})
+ for _, keyspace := range mergeTargetKg.Keyspaces {
+ keyspaces[keyspace] = struct{}{}
+ }
+ // Delete the keyspace groups in merge list and move the keyspaces in it to the target keyspace group.
+ for _, kgID := range mergeList {
+ kg := groups[kgID]
+ for _, keyspace := range kg.Keyspaces {
+ keyspaces[keyspace] = struct{}{}
+ }
+ if err := m.store.DeleteKeyspaceGroup(txn, kg.ID); err != nil {
+ return err
+ }
+ }
+ mergedKeyspaces := make([]uint32, 0, len(keyspaces))
+ for keyspace := range keyspaces {
+ mergedKeyspaces = append(mergedKeyspaces, keyspace)
+ }
+ sort.Slice(mergedKeyspaces, func(i, j int) bool {
+ return mergedKeyspaces[i] < mergedKeyspaces[j]
+ })
+ mergeTargetKg.Keyspaces = mergedKeyspaces
+ // Update the merge state of the target keyspace group.
+ mergeTargetKg.MergeState = &endpoint.MergeState{
+ MergeList: mergeList,
+ }
+ return m.store.SaveKeyspaceGroup(txn, mergeTargetKg)
+ }); err != nil {
+ return err
+ }
+ // Update the keyspace group cache.
+ m.groups[endpoint.StringUserKind(mergeTargetKg.UserKind)].Put(mergeTargetKg)
+ for _, kgID := range mergeList {
+ kg := groups[kgID]
+ m.groups[endpoint.StringUserKind(kg.UserKind)].Remove(kgID)
+ }
+ return nil
+}
+
+// FinishMergeKeyspaceByID finishes the merging keyspace group by the merge target ID.
+func (m *GroupManager) FinishMergeKeyspaceByID(mergeTargetID uint32) error {
+ var mergeTargetKg *endpoint.KeyspaceGroup
+ m.Lock()
+ defer m.Unlock()
+ if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) (err error) {
+ // Load the merge target keyspace group first.
+ mergeTargetKg, err = m.store.LoadKeyspaceGroup(txn, mergeTargetID)
+ if err != nil {
+ return err
+ }
+ if mergeTargetKg == nil {
+ return ErrKeyspaceGroupNotExists
+ }
+ // Check if it's in the merging state.
+ if !mergeTargetKg.IsMergeTarget() {
+ return ErrKeyspaceGroupNotInMerging
+ }
+ // Make sure all merging keyspace groups are deleted.
+ for _, kgID := range mergeTargetKg.MergeState.MergeList {
+ kg, err := m.store.LoadKeyspaceGroup(txn, kgID)
+ if err != nil {
+ return err
+ }
+ if kg != nil {
+ return ErrKeyspaceGroupNotInMerging
+ }
+ }
+ mergeTargetKg.MergeState = nil
+ return m.store.SaveKeyspaceGroup(txn, mergeTargetKg)
+ }); err != nil {
+ return err
+ }
+ // Update the keyspace group cache.
+ m.groups[endpoint.StringUserKind(mergeTargetKg.UserKind)].Put(mergeTargetKg)
+ return nil
+}

pkg/keyspace/tso_keyspace_group_test.go

@@ -84,11 +84,12 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() {
  re.NoError(err)
  re.Len(kgs, 2)
  // get the default keyspace group
- kg, err := suite.kgm.GetKeyspaceGroupByID(0)
+ kg, err := suite.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID)
  re.NoError(err)
  re.Equal(uint32(0), kg.ID)
  re.Equal(endpoint.Basic.String(), kg.UserKind)
  re.False(kg.IsSplitting())
+ // get the keyspace group 3
  kg, err = suite.kgm.GetKeyspaceGroupByID(3)
  re.NoError(err)
  re.Equal(uint32(3), kg.ID)
@@ -320,3 +321,80 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() {
  err = suite.kgm.SplitKeyspaceGroupByID(2, 5, []uint32{111, 222, 444})
  re.ErrorIs(err, ErrKeyspaceNotInKeyspaceGroup)
 }
+
+func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupMerge() {
+ re := suite.Require()
+
+ keyspaceGroups := []*endpoint.KeyspaceGroup{
+ {
+ ID: uint32(1),
+ UserKind: endpoint.Basic.String(),
+ Keyspaces: []uint32{111, 222, 333},
+ Members: make([]endpoint.KeyspaceGroupMember, utils.KeyspaceGroupDefaultReplicaCount),
+ },
+ {
+ ID: uint32(3),
+ UserKind: endpoint.Basic.String(),
+ Keyspaces: []uint32{444, 555},
+ },
+ }
+ err := suite.kgm.CreateKeyspaceGroups(keyspaceGroups)
+ re.NoError(err)
+ // split the keyspace group 1 to 2
+ err = suite.kgm.SplitKeyspaceGroupByID(1, 2, []uint32{333})
+ re.NoError(err)
+ // finish the split of the keyspace group 2
+ err = suite.kgm.FinishSplitKeyspaceByID(2)
+ re.NoError(err)
+ // check the keyspace group 1 and 2
+ kg1, err := suite.kgm.GetKeyspaceGroupByID(1)
+ re.NoError(err)
+ re.Equal(uint32(1), kg1.ID)
+ re.Equal([]uint32{111, 222}, kg1.Keyspaces)
+ re.False(kg1.IsSplitting())
+ re.False(kg1.IsMerging())
+ kg2, err := suite.kgm.GetKeyspaceGroupByID(2)
+ re.NoError(err)
+ re.Equal(uint32(2), kg2.ID)
+ re.Equal([]uint32{333}, kg2.Keyspaces)
+ re.False(kg2.IsSplitting())
+ re.False(kg2.IsMerging())
+ re.Equal(kg1.UserKind, kg2.UserKind)
+ re.Equal(kg1.Members, kg2.Members)
+ // merge the keyspace group 2 and 3 back into 1
+ err = suite.kgm.MergeKeyspaceGroups(1, []uint32{2, 3})
+ re.NoError(err)
+ // check the keyspace group 2 and 3
+ kg2, err = suite.kgm.GetKeyspaceGroupByID(2)
+ re.NoError(err)
+ re.Nil(kg2)
+ kg3, err := suite.kgm.GetKeyspaceGroupByID(3)
+ re.NoError(err)
+ re.Nil(kg3)
+ // check the keyspace group 1
+ kg1, err = suite.kgm.GetKeyspaceGroupByID(1)
+ re.NoError(err)
+ re.Equal(uint32(1), kg1.ID)
+ re.Equal([]uint32{111, 222, 333, 444, 555}, kg1.Keyspaces)
+ re.False(kg1.IsSplitting())
+ re.True(kg1.IsMerging())
+ // finish the merging
+ err = suite.kgm.FinishMergeKeyspaceByID(1)
+ re.NoError(err)
+ kg1, err = suite.kgm.GetKeyspaceGroupByID(1)
+ re.NoError(err)
+ re.Equal(uint32(1), kg1.ID)
+ re.Equal([]uint32{111, 222, 333, 444, 555}, kg1.Keyspaces)
+ re.False(kg1.IsSplitting())
+ re.False(kg1.IsMerging())
+
+ // merge a non-existing keyspace group
+ err = suite.kgm.MergeKeyspaceGroups(4, []uint32{5})
+ re.ErrorIs(err, ErrKeyspaceGroupNotExists)
+ // merge with the number of keyspace groups exceeds the limit
+ err = suite.kgm.MergeKeyspaceGroups(1, make([]uint32, maxEtcdTxnOps/2))
+ re.ErrorIs(err, ErrExceedMaxEtcdTxnOps)
+ // merge the default keyspace group
+ err = suite.kgm.MergeKeyspaceGroups(1, []uint32{utils.DefaultKeyspaceGroupID})
+ re.ErrorIs(err, ErrModifyDefaultKeyspaceGroup)
+}