Address the comments

Signed-off-by: JmPotato <[email protected]>
tikv · Jun 14, 2023 · 90ae755 · 90ae755
1 parent 9a43229
commit 90ae755
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 5 deletions.
diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go
@@ -49,7 +49,7 @@ const (
 	UserKindKey = "user_kind"
 	// TSOKeyspaceGroupIDKey is the key for tso keyspace group id in keyspace config.
 	TSOKeyspaceGroupIDKey = "tso_keyspace_group_id"
-	// maxEtcdTxnOps is the batch size for operating etcd. The limit of etcd txn op is 128.
+	// 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
 	maxEtcdTxnOps = 120

diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go
@@ -766,7 +766,11 @@ func (m *GroupManager) MergeKeyspaceGroups(mergeTargetID uint32, mergeList []uin
 	if mergeListNum == 0 {
 		return nil
 	}
-	if mergeListNum > maxEtcdTxnOps {
+	// 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
 	}
 	var (
@@ -796,7 +800,7 @@ func (m *GroupManager) MergeKeyspaceGroups(mergeTargetID uint32, mergeList []uin
 			groups[kgID] = kg
 		}
 		mergeTargetKg = groups[mergeTargetID]
-		keyspaces := make(map[uint32]struct{}, len(mergeTargetKg.Keyspaces))
+		keyspaces := make(map[uint32]struct{})
 		for _, keyspace := range mergeTargetKg.Keyspaces {
 			keyspaces[keyspace] = struct{}{}
 		}

diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go
@@ -387,4 +387,11 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupMerge() {
 	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)
 }
diff --git a/server/apiv2/handlers/tso_keyspace_group.go b/server/apiv2/handlers/tso_keyspace_group.go
@@ -207,12 +207,12 @@ func SplitKeyspaceGroupByID(c *gin.Context) {
 		patrolKeyspaceAssignmentState.patrolled = true
 	}
 	patrolKeyspaceAssignmentState.Unlock()
-	// Split keyspace group.
 	groupManager := svr.GetKeyspaceGroupManager()
 	if groupManager == nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, groupManagerUninitializedErr)
 		return
 	}
+	// Split keyspace group.
 	err = groupManager.SplitKeyspaceGroupByID(id, splitParams.NewID, splitParams.Keyspaces)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -269,12 +269,12 @@ func MergeKeyspaceGroups(c *gin.Context) {
 	}
 
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
-	// Split keyspace group.
 	groupManager := svr.GetKeyspaceGroupManager()
 	if groupManager == nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, groupManagerUninitializedErr)
 		return
 	}
+	// Merge keyspace group.
 	err = groupManager.MergeKeyspaceGroups(id, mergeParams.MergeList)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())