tools: add merge commands for pd-ctl (tikv#6675)

ref tikv#6589 Signed-off-by: Ryan Leung <[email protected]> Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
lhy1024 · Jun 26, 2023 · 4d26c8b · 4d26c8b
1 parent b4b1fcc
commit 4d26c8b
Show file tree

Hide file tree

Showing 2 changed files with 165 additions and 0 deletions.
diff --git 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
@@ -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,35 @@ func newSplitKeyspaceGroupCommand() *cobra.Command {
 	return r
 }
 
+func newFinishSplitKeyspaceGroupCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:    "finish-split <keyspace_group_id>",
+		Short:  "finish split the keyspace group with the given ID",
+		Run:    finishSplitKeyspaceGroupCommandFunc,
+		Hidden: true,
+	}
+	return r
+}
+
+func newMergeKeyspaceGroupCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "merge <target_keyspace_group_id> [<keyspace_group_id>]",
+		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 <keyspace_group_id>",
+		Short:  "finish merge the keyspace group with the given ID",
+		Run:    finishMergeKeyspaceGroupCommandFunc,
+		Hidden: true,
+	}
+	return r
+}
+
 func newSetNodesKeyspaceGroupCommand() *cobra.Command {
 	r := &cobra.Command{
 		Use:   "set-node <keyspace_group_id> <tso_node_addr> [<tso_node_addr>...]",
@@ -108,6 +140,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()