diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index b3c81144f01e..feda34f48b10 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -656,7 +656,13 @@ func (manager *Manager) allocID() (uint32, error) { } // PatrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. -func (manager *Manager) PatrolKeyspaceAssignment() error { +func (manager *Manager) PatrolKeyspaceAssignment(startKeyspaceID, endKeyspaceID uint32) error { + if startKeyspaceID > manager.nextPatrolStartID { + manager.nextPatrolStartID = startKeyspaceID + } + if endKeyspaceID != 0 && endKeyspaceID < manager.nextPatrolStartID { + return nil + } var ( // Some statistics info. start = time.Now() @@ -675,6 +681,8 @@ func (manager *Manager) PatrolKeyspaceAssignment() error { zap.Uint64("patrolled-keyspace-count", patrolledKeyspaceCount), zap.Uint64("assigned-keyspace-count", assignedKeyspaceCount), zap.Int("batch-size", maxEtcdTxnOps), + zap.Uint32("start-keyspace-id", startKeyspaceID), + zap.Uint32("end-keyspace-id", endKeyspaceID), zap.Uint32("current-start-id", currentStartID), zap.Uint32("next-start-id", nextStartID), ) @@ -706,9 +714,9 @@ func (manager *Manager) PatrolKeyspaceAssignment() error { currentStartID = keyspaces[0].GetId() nextStartID = keyspaces[keyspaceNum-1].GetId() + 1 } - // If there are less than `maxEtcdTxnOps` keyspaces, - // we have reached the end of the keyspace list. - moreToPatrol = keyspaceNum == maxEtcdTxnOps + // If there are less than `maxEtcdTxnOps` keyspaces or the next start ID reaches the end, + // there is no need to patrol again. + moreToPatrol = keyspaceNum == maxEtcdTxnOps && endKeyspaceID != 0 && nextStartID < endKeyspaceID var ( assigned = false keyspaceIDsToUnlock = make([]uint32, 0, keyspaceNum) @@ -722,6 +730,10 @@ func (manager *Manager) PatrolKeyspaceAssignment() error { if ks == nil { continue } + if endKeyspaceID != 0 && ks.Id > endKeyspaceID { + moreToPatrol = false + break + } patrolledKeyspaceCount++ manager.metaLock.Lock(ks.Id) if ks.Config == nil { @@ -744,6 +756,8 @@ func (manager *Manager) PatrolKeyspaceAssignment() error { if err != nil { log.Error("[keyspace] failed to save keyspace meta during patrol", zap.Int("batch-size", maxEtcdTxnOps), + zap.Uint32("start-keyspace-id", startKeyspaceID), + zap.Uint32("end-keyspace-id", endKeyspaceID), zap.Uint32("current-start-id", currentStartID), zap.Uint32("next-start-id", nextStartID), zap.Uint32("keyspace-id", ks.Id), zap.Error(err)) @@ -756,6 +770,8 @@ func (manager *Manager) PatrolKeyspaceAssignment() error { if err != nil { log.Error("[keyspace] failed to save default keyspace group meta during patrol", zap.Int("batch-size", maxEtcdTxnOps), + zap.Uint32("start-keyspace-id", startKeyspaceID), + zap.Uint32("end-keyspace-id", endKeyspaceID), zap.Uint32("current-start-id", currentStartID), zap.Uint32("next-start-id", nextStartID), zap.Error(err)) return err diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go index 45c4bc90be23..bbaedef2b542 100644 --- a/pkg/keyspace/keyspace_test.go +++ b/pkg/keyspace/keyspace_test.go @@ -393,7 +393,7 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignment() { re.NotNil(defaultKeyspaceGroup) re.NotContains(defaultKeyspaceGroup.Keyspaces, uint32(111)) // Patrol the keyspace assignment. - err = suite.manager.PatrolKeyspaceAssignment() + err = suite.manager.PatrolKeyspaceAssignment(0, 0) re.NoError(err) // Check if the keyspace is attached to the default group. defaultKeyspaceGroup, err = suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID) @@ -424,7 +424,7 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() { re.NotContains(defaultKeyspaceGroup.Keyspaces, uint32(i)) } // Patrol the keyspace assignment. - err = suite.manager.PatrolKeyspaceAssignment() + err = suite.manager.PatrolKeyspaceAssignment(0, 0) re.NoError(err) // Check if all the keyspaces are attached to the default group. defaultKeyspaceGroup, err = suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID) @@ -435,6 +435,43 @@ func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentInBatch() { } } +func (suite *keyspaceTestSuite) TestPatrolKeyspaceAssignmentWithRange() { + re := suite.Require() + // Create some keyspaces without any keyspace group. + for i := 1; i < maxEtcdTxnOps*2+1; i++ { + now := time.Now().Unix() + err := suite.manager.saveNewKeyspace(&keyspacepb.KeyspaceMeta{ + Id: uint32(i), + Name: strconv.Itoa(i), + State: keyspacepb.KeyspaceState_ENABLED, + CreatedAt: now, + StateChangedAt: now, + }) + re.NoError(err) + } + // Check if all the keyspaces are not attached to the default group. + defaultKeyspaceGroup, err := suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID) + re.NoError(err) + re.NotNil(defaultKeyspaceGroup) + for i := 1; i < maxEtcdTxnOps*2+1; i++ { + re.NotContains(defaultKeyspaceGroup.Keyspaces, uint32(i)) + } + // Patrol the keyspace assignment with range [10, 20] + err = suite.manager.PatrolKeyspaceAssignment(10, 20) + re.NoError(err) + // Check if only the keyspaces within the range are attached to the default group. + defaultKeyspaceGroup, err = suite.manager.kgm.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID) + re.NoError(err) + re.NotNil(defaultKeyspaceGroup) + for i := 1; i < maxEtcdTxnOps*2+1; i++ { + if i >= 10 && i <= 20 { + re.Contains(defaultKeyspaceGroup.Keyspaces, uint32(i)) + } else { + re.NotContains(defaultKeyspaceGroup.Keyspaces, uint32(i)) + } + } +} + // Benchmark the keyspace assignment patrol. func BenchmarkPatrolKeyspaceAssignment1000(b *testing.B) { benchmarkPatrolKeyspaceAssignmentN(1000, b) @@ -471,7 +508,7 @@ func benchmarkPatrolKeyspaceAssignmentN( // Benchmark the keyspace assignment patrol. b.ResetTimer() for i := 0; i < b.N; i++ { - err := suite.manager.PatrolKeyspaceAssignment() + err := suite.manager.PatrolKeyspaceAssignment(0, 0) re.NoError(err) } b.StopTimer() diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index c5e78d6b7979..20d1307a9472 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -509,7 +509,10 @@ func (m *GroupManager) UpdateKeyspaceGroup(oldGroupID, newGroupID string, oldUse // SplitKeyspaceGroupByID splits the keyspace group by ID into a new keyspace group with the given new ID. // And the keyspaces in the old keyspace group will be moved to the new keyspace group. -func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint32, keyspaces []uint32) error { +func (m *GroupManager) SplitKeyspaceGroupByID( + splitSourceID, splitTargetID uint32, + keyspaces []uint32, keyspaceIDRange ...uint32, +) error { var splitSourceKg, splitTargetKg *endpoint.KeyspaceGroup m.Lock() defer m.Unlock() @@ -542,34 +545,17 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3 if splitTargetKg != nil { return ErrKeyspaceGroupExists } - keyspaceNum := len(keyspaces) - sourceKeyspaceNum := len(splitSourceKg.Keyspaces) - // Check if the keyspaces are all in the old keyspace group. - if keyspaceNum == 0 || keyspaceNum > sourceKeyspaceNum { - return ErrKeyspaceNotInKeyspaceGroup + var startKeyspaceID, endKeyspaceID uint32 + if len(keyspaceIDRange) >= 2 { + startKeyspaceID, endKeyspaceID = keyspaceIDRange[0], keyspaceIDRange[1] } - var ( - oldKeyspaceMap = make(map[uint32]struct{}, sourceKeyspaceNum) - newKeyspaceMap = make(map[uint32]struct{}, keyspaceNum) - ) - for _, keyspace := range splitSourceKg.Keyspaces { - oldKeyspaceMap[keyspace] = struct{}{} - } - for _, keyspace := range keyspaces { - if _, ok := oldKeyspaceMap[keyspace]; !ok { - return ErrKeyspaceNotInKeyspaceGroup - } - newKeyspaceMap[keyspace] = struct{}{} - } - // Get the split keyspace group for the old keyspace group. - splitKeyspaces := make([]uint32, 0, sourceKeyspaceNum-keyspaceNum) - for _, keyspace := range splitSourceKg.Keyspaces { - if _, ok := newKeyspaceMap[keyspace]; !ok { - splitKeyspaces = append(splitKeyspaces, keyspace) - } + splitSourceKeyspaces, splitTargetKeyspaces, err := buildSplitKeyspaces( + splitSourceKg.Keyspaces, keyspaces, startKeyspaceID, endKeyspaceID) + if err != nil { + return err } // Update the old keyspace group. - splitSourceKg.Keyspaces = splitKeyspaces + splitSourceKg.Keyspaces = splitSourceKeyspaces splitSourceKg.SplitState = &endpoint.SplitState{ SplitSource: splitSourceKg.ID, } @@ -581,7 +567,7 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3 // Keep the same user kind and members as the old keyspace group. UserKind: splitSourceKg.UserKind, Members: splitSourceKg.Members, - Keyspaces: keyspaces, + Keyspaces: splitTargetKeyspaces, SplitState: &endpoint.SplitState{ SplitSource: splitSourceKg.ID, }, @@ -597,6 +583,64 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3 return nil } +func buildSplitKeyspaces( + // `old` is the original keyspace list which will be split out, + // `new` is the keyspace list which will be split from the old keyspace list. + old, new []uint32, + startKeyspaceID, endKeyspaceID uint32, +) ([]uint32, []uint32, error) { + oldNum, newNum := len(old), len(new) + // Split according to the new keyspace list. + if newNum != 0 { + if newNum > oldNum { + return nil, nil, ErrKeyspaceNotInKeyspaceGroup + } + var ( + oldKeyspaceMap = make(map[uint32]struct{}, oldNum) + newKeyspaceMap = make(map[uint32]struct{}, newNum) + ) + for _, keyspace := range old { + oldKeyspaceMap[keyspace] = struct{}{} + } + for _, keyspace := range new { + if _, ok := oldKeyspaceMap[keyspace]; !ok { + return nil, nil, ErrKeyspaceNotInKeyspaceGroup + } + newKeyspaceMap[keyspace] = struct{}{} + } + // Get the split keyspace list for the old keyspace group. + oldSplit := make([]uint32, 0, oldNum-newNum) + for _, keyspace := range old { + if _, ok := newKeyspaceMap[keyspace]; !ok { + oldSplit = append(oldSplit, keyspace) + } + } + return oldSplit, new, nil + } + // Split according to the start and end keyspace ID. + if startKeyspaceID == 0 && endKeyspaceID == 0 { + return nil, nil, ErrKeyspaceNotInKeyspaceGroup + } + var ( + newSplit = make([]uint32, 0, oldNum) + newKeyspaceMap = make(map[uint32]struct{}, newNum) + ) + for _, keyspace := range old { + if startKeyspaceID <= keyspace && keyspace <= endKeyspaceID { + newSplit = append(newSplit, keyspace) + newKeyspaceMap[keyspace] = struct{}{} + } + } + // Get the split keyspace list for the old keyspace group. + oldSplit := make([]uint32, 0, oldNum-len(newSplit)) + for _, keyspace := range old { + if _, ok := newKeyspaceMap[keyspace]; !ok { + oldSplit = append(oldSplit, keyspace) + } + } + return oldSplit, newSplit, nil +} + // FinishSplitKeyspaceByID finishes the split keyspace group by the split target ID. func (m *GroupManager) FinishSplitKeyspaceByID(splitTargetID uint32) error { var splitTargetKg, splitSourceKg *endpoint.KeyspaceGroup diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index 60f2793b8bb7..42c8918e78b2 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -20,6 +20,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "github.com/tikv/pd/pkg/mcs/utils" "github.com/tikv/pd/pkg/mock/mockcluster" @@ -322,6 +323,57 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() { re.ErrorIs(err, ErrKeyspaceNotInKeyspaceGroup) } +func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplitRange() { + re := suite.Require() + + keyspaceGroups := []*endpoint.KeyspaceGroup{ + { + ID: uint32(1), + UserKind: endpoint.Basic.String(), + }, + { + ID: uint32(2), + UserKind: endpoint.Standard.String(), + Keyspaces: []uint32{111, 333, 444, 555, 666}, + Members: make([]endpoint.KeyspaceGroupMember, utils.DefaultKeyspaceGroupReplicaCount), + }, + } + err := suite.kgm.CreateKeyspaceGroups(keyspaceGroups) + re.NoError(err) + // split the keyspace group 2 to 4 with keyspace range [222, 555] + err = suite.kgm.SplitKeyspaceGroupByID(2, 4, nil, 222, 555) + re.NoError(err) + kg2, err := suite.kgm.GetKeyspaceGroupByID(2) + re.NoError(err) + re.Equal(uint32(2), kg2.ID) + re.Equal([]uint32{111, 666}, kg2.Keyspaces) + re.True(kg2.IsSplitSource()) + re.Equal(kg2.ID, kg2.SplitSource()) + kg4, err := suite.kgm.GetKeyspaceGroupByID(4) + re.NoError(err) + re.Equal(uint32(4), kg4.ID) + re.Equal([]uint32{333, 444, 555}, kg4.Keyspaces) + re.True(kg4.IsSplitTarget()) + re.Equal(kg2.ID, kg4.SplitSource()) + re.Equal(kg2.UserKind, kg4.UserKind) + re.Equal(kg2.Members, kg4.Members) + // finish the split of keyspace group 4 + err = suite.kgm.FinishSplitKeyspaceByID(4) + re.NoError(err) + kg2, err = suite.kgm.GetKeyspaceGroupByID(2) + re.NoError(err) + re.Equal(uint32(2), kg2.ID) + re.Equal([]uint32{111, 666}, kg2.Keyspaces) + re.False(kg2.IsSplitting()) + kg4, err = suite.kgm.GetKeyspaceGroupByID(4) + re.NoError(err) + re.Equal(uint32(4), kg4.ID) + re.Equal([]uint32{333, 444, 555}, kg4.Keyspaces) + re.False(kg4.IsSplitting()) + re.Equal(kg2.UserKind, kg4.UserKind) + re.Equal(kg2.Members, kg4.Members) +} + func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupMerge() { re := suite.Require() @@ -398,3 +450,69 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupMerge() { err = suite.kgm.MergeKeyspaceGroups(1, []uint32{utils.DefaultKeyspaceGroupID}) re.ErrorIs(err, ErrModifyDefaultKeyspaceGroup) } + +func TestBuildSplitKeyspaces(t *testing.T) { + re := require.New(t) + testCases := []struct { + old []uint32 + new []uint32 + startKeyspaceID uint32 + endKeyspaceID uint32 + expectedOld []uint32 + expectedNew []uint32 + err error + }{ + { + old: []uint32{1, 2, 3, 4, 5}, + new: []uint32{1, 2, 3, 4, 5}, + expectedOld: []uint32{}, + expectedNew: []uint32{1, 2, 3, 4, 5}, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + new: []uint32{1}, + expectedOld: []uint32{2, 3, 4, 5}, + expectedNew: []uint32{1}, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + new: []uint32{6}, + err: ErrKeyspaceNotInKeyspaceGroup, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + startKeyspaceID: 2, + endKeyspaceID: 4, + expectedOld: []uint32{1, 5}, + expectedNew: []uint32{2, 3, 4}, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + startKeyspaceID: 2, + endKeyspaceID: 6, + expectedOld: []uint32{1}, + expectedNew: []uint32{2, 3, 4, 5}, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + startKeyspaceID: 0, + endKeyspaceID: 6, + expectedOld: []uint32{}, + expectedNew: []uint32{1, 2, 3, 4, 5}, + }, + { + old: []uint32{1, 2, 3, 4, 5}, + err: ErrKeyspaceNotInKeyspaceGroup, + }, + } + for idx, testCase := range testCases { + old, new, err := buildSplitKeyspaces(testCase.old, testCase.new, testCase.startKeyspaceID, testCase.endKeyspaceID) + if testCase.err != nil { + re.ErrorIs(testCase.err, err, "test case %d", idx) + } else { + re.NoError(err, "test case %d", idx) + re.Equal(testCase.expectedOld, old, "test case %d", idx) + re.Equal(testCase.expectedNew, new, "test case %d", idx) + } + } +} diff --git a/server/apiv2/handlers/tso_keyspace_group.go b/server/apiv2/handlers/tso_keyspace_group.go index 2d16f6ac360c..e7c06bf608b8 100644 --- a/server/apiv2/handlers/tso_keyspace_group.go +++ b/server/apiv2/handlers/tso_keyspace_group.go @@ -161,6 +161,9 @@ func DeleteKeyspaceGroupByID(c *gin.Context) { type SplitKeyspaceGroupByIDParams struct { NewID uint32 `json:"new-id"` Keyspaces []uint32 `json:"keyspaces"` + // StartKeyspaceID and EndKeyspaceID are used to indicate the range of keyspaces to be split. + StartKeyspaceID uint32 `json:"start-keyspace-id"` + EndKeyspaceID uint32 `json:"end-keyspace-id"` } var patrolKeyspaceAssignmentState struct { @@ -186,10 +189,15 @@ func SplitKeyspaceGroupByID(c *gin.Context) { c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id") return } - if len(splitParams.Keyspaces) == 0 { + if len(splitParams.Keyspaces) == 0 && splitParams.StartKeyspaceID == 0 && splitParams.EndKeyspaceID == 0 { c.AbortWithStatusJSON(http.StatusBadRequest, "invalid empty keyspaces") return } + if splitParams.StartKeyspaceID < utils.DefaultKeyspaceID || + splitParams.StartKeyspaceID > splitParams.EndKeyspaceID { + c.AbortWithStatusJSON(http.StatusBadRequest, "invalid start/end keyspace id") + return + } svr := c.MustGet(middlewares.ServerContextKey).(*server.Server) patrolKeyspaceAssignmentState.Lock() @@ -200,7 +208,7 @@ func SplitKeyspaceGroupByID(c *gin.Context) { c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr) return } - err = manager.PatrolKeyspaceAssignment() + err = manager.PatrolKeyspaceAssignment(splitParams.StartKeyspaceID, splitParams.EndKeyspaceID) if err != nil { c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) patrolKeyspaceAssignmentState.Unlock() @@ -215,7 +223,9 @@ func SplitKeyspaceGroupByID(c *gin.Context) { return } // Split keyspace group. - err = groupManager.SplitKeyspaceGroupByID(id, splitParams.NewID, splitParams.Keyspaces) + err = groupManager.SplitKeyspaceGroupByID( + id, splitParams.NewID, + splitParams.Keyspaces, splitParams.StartKeyspaceID, splitParams.EndKeyspaceID) if err != nil { c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) return