diff --git a/tests/pdctl/keyspace/keyspace_group_test.go b/tests/pdctl/keyspace/keyspace_group_test.go index 0c44700f48e..80f9eaa8420 100644 --- a/tests/pdctl/keyspace/keyspace_group_test.go +++ b/tests/pdctl/keyspace/keyspace_group_test.go @@ -277,3 +277,76 @@ func TestSetNodeAndPriorityKeyspaceGroup(t *testing.T) { re.NoError(err) re.Contains(string(output), "Failed to parse the priority") } + +func TestMergeKeyspaceGroup(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes", `return(true)`)) + re.NoError(failpoint.Enable("github.com/tikv/pd/server/delayStartServerLoop", `return(true)`)) + keyspaces := make([]string, 0) + // we test the case which exceed the default max txn ops limit in etcd, which is 128. + for i := 0; i < 129; i++ { + keyspaces = append(keyspaces, fmt.Sprintf("keyspace_%d", i)) + } + tc, err := tests.NewTestAPICluster(ctx, 3, func(conf *config.Config, serverName string) { + conf.Keyspace.PreAlloc = keyspaces + }) + re.NoError(err) + err = tc.RunInitialServers() + re.NoError(err) + pdAddr := tc.GetConfig().GetClientURL() + + _, tsoServerCleanup1, err := tests.StartSingleTSOTestServer(ctx, re, pdAddr, tempurl.Alloc()) + defer tsoServerCleanup1() + re.NoError(err) + _, tsoServerCleanup2, err := tests.StartSingleTSOTestServer(ctx, re, pdAddr, tempurl.Alloc()) + defer tsoServerCleanup2() + re.NoError(err) + cmd := pdctlCmd.GetRootCmd() + + tc.WaitLeader() + leaderServer := tc.GetServer(tc.GetLeader()) + re.NoError(leaderServer.BootstrapCluster()) + + // split keyspace group. + testutil.Eventually(re, func() bool { + args := []string{"-u", pdAddr, "keyspace-group", "split", "0", "1", "2"} + output, err := pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + return strings.Contains(string(output), "Success") + }) + + args := []string{"-u", pdAddr, "keyspace-group", "finish-split", "0"} + output, err := pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + args = []string{"-u", pdAddr, "keyspace-group", "finish-split", "1"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + + // merge keyspace group. + testutil.Eventually(re, func() bool { + args := []string{"-u", pdAddr, "keyspace-group", "merge", "0", "1"} + output, err := pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + return strings.Contains(string(output), "Success") + }) + + args = []string{"-u", pdAddr, "keyspace-group", "finish-merge", "0"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + args = []string{"-u", pdAddr, "keyspace-group", "0"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + var keyspaceGroup endpoint.KeyspaceGroup + err = json.Unmarshal(output, &keyspaceGroup) + re.NoError(err) + re.Len(keyspaceGroup.Keyspaces, 130) + re.Nil(keyspaceGroup.MergeState) + + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes")) + re.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop")) +} diff --git a/tools/pd-ctl/pdctl/command/keyspace_group_command.go b/tools/pd-ctl/pdctl/command/keyspace_group_command.go index 662a4aa157e..3485cf39020 100644 --- a/tools/pd-ctl/pdctl/command/keyspace_group_command.go +++ b/tools/pd-ctl/pdctl/command/keyspace_group_command.go @@ -33,6 +33,9 @@ func NewKeyspaceGroupCommand() *cobra.Command { Run: showKeyspaceGroupCommandFunc, } cmd.AddCommand(newSplitKeyspaceGroupCommand()) + cmd.AddCommand(newFinishSplitKeyspaceGroupCommand()) + cmd.AddCommand(newMergeKeyspaceGroupCommand()) + cmd.AddCommand(newFinishMergeKeyspaceGroupCommand()) cmd.AddCommand(newSetNodesKeyspaceGroupCommand()) cmd.AddCommand(newSetPriorityKeyspaceGroupCommand()) return cmd @@ -47,6 +50,33 @@ func newSplitKeyspaceGroupCommand() *cobra.Command { return r } +func newFinishSplitKeyspaceGroupCommand() *cobra.Command { + r := &cobra.Command{ + Use: "finish-split ", + Short: "finish split the keyspace group with the given ID", + Run: finishSplitKeyspaceGroupCommandFunc, + } + return r +} + +func newMergeKeyspaceGroupCommand() *cobra.Command { + r := &cobra.Command{ + Use: "merge []", + Short: "merge the keyspace group with the given IDs into the target one", + Run: mergeKeyspaceGroupCommandFunc, + } + return r +} + +func newFinishMergeKeyspaceGroupCommand() *cobra.Command { + r := &cobra.Command{ + Use: "finish-merge ", + Short: "finish merge the keyspace group with the given ID", + Run: finishMergeKeyspaceGroupCommandFunc, + } + return r +} + func newSetNodesKeyspaceGroupCommand() *cobra.Command { r := &cobra.Command{ Use: "set-node [...]", @@ -108,6 +138,66 @@ func splitKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { }) } +func finishSplitKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { + if len(args) < 1 { + cmd.Usage() + return + } + _, err := strconv.ParseUint(args[0], 10, 32) + if err != nil { + cmd.Printf("Failed to parse the keyspace group ID: %s\n", err) + return + } + _, err = doRequest(cmd, fmt.Sprintf("%s/%s/split", keyspaceGroupsPrefix, args[0]), http.MethodDelete, http.Header{}) + if err != nil { + cmd.Println(err) + return + } + cmd.Println("Success!") +} + +func mergeKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { + if len(args) < 2 { + cmd.Usage() + return + } + _, err := strconv.ParseUint(args[0], 10, 32) + if err != nil { + cmd.Printf("Failed to parse the target keyspace group ID: %s\n", err) + return + } + groups := make([]uint32, 0, len(args)-1) + for _, arg := range args[1:] { + id, err := strconv.ParseUint(arg, 10, 32) + if err != nil { + cmd.Printf("Failed to parse the keyspace ID: %s\n", err) + return + } + groups = append(groups, uint32(id)) + } + postJSON(cmd, fmt.Sprintf("%s/%s/merge", keyspaceGroupsPrefix, args[0]), map[string]interface{}{ + "merge-list": groups, + }) +} + +func finishMergeKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { + if len(args) < 1 { + cmd.Usage() + return + } + _, err := strconv.ParseUint(args[0], 10, 32) + if err != nil { + cmd.Printf("Failed to parse the keyspace group ID: %s\n", err) + return + } + _, err = doRequest(cmd, fmt.Sprintf("%s/%s/merge", keyspaceGroupsPrefix, args[0]), http.MethodDelete, http.Header{}) + if err != nil { + cmd.Println(err) + return + } + cmd.Println("Success!") +} + func setNodesKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { if len(args) < 2 { cmd.Usage()