From 7c6cc4f2c505f892fc58ab7f6806f40ddaea94a0 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 29 Jun 2023 10:16:43 +0800 Subject: [PATCH] tools: support querying keyspace group by state (#6706) ref tikv/pd#5895 Signed-off-by: Ryan Leung Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com> --- server/apiv2/handlers/tso_keyspace_group.go | 25 ++++- tests/pdctl/keyspace/keyspace_group_test.go | 100 ++++++++++++++++-- .../pdctl/command/keyspace_group_command.go | 72 +++++++++++-- 3 files changed, 184 insertions(+), 13 deletions(-) diff --git a/server/apiv2/handlers/tso_keyspace_group.go b/server/apiv2/handlers/tso_keyspace_group.go index e7c06bf608b..5de8fd6a9cc 100644 --- a/server/apiv2/handlers/tso_keyspace_group.go +++ b/server/apiv2/handlers/tso_keyspace_group.go @@ -17,6 +17,7 @@ package handlers import ( "net/http" "strconv" + "strings" "sync" "github.com/gin-gonic/gin" @@ -108,8 +109,30 @@ func GetKeyspaceGroups(c *gin.Context) { c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) return } + var kgs []*endpoint.KeyspaceGroup + state, set := c.GetQuery("state") + if set { + state := strings.ToLower(state) + switch state { + case "merge": + for _, keyspaceGroup := range keyspaceGroups { + if keyspaceGroup.MergeState != nil { + kgs = append(kgs, keyspaceGroup) + } + } + case "split": + for _, keyspaceGroup := range keyspaceGroups { + if keyspaceGroup.SplitState != nil { + kgs = append(kgs, keyspaceGroup) + } + } + default: + } + } else { + kgs = keyspaceGroups + } - c.IndentedJSON(http.StatusOK, keyspaceGroups) + c.IndentedJSON(http.StatusOK, kgs) } // GetKeyspaceGroupByID gets keyspace group by ID. diff --git a/tests/pdctl/keyspace/keyspace_group_test.go b/tests/pdctl/keyspace/keyspace_group_test.go index 8715df7f3da..e71c1e0f0df 100644 --- a/tests/pdctl/keyspace/keyspace_group_test.go +++ b/tests/pdctl/keyspace/keyspace_group_test.go @@ -297,7 +297,7 @@ func TestMergeKeyspaceGroup(t *testing.T) { 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) { + tc, err := tests.NewTestAPICluster(ctx, 1, func(conf *config.Config, serverName string) { conf.Keyspace.PreAlloc = keyspaces }) re.NoError(err) @@ -325,14 +325,10 @@ func TestMergeKeyspaceGroup(t *testing.T) { return strings.Contains(string(output), "Success") }) - args := []string{"-u", pdAddr, "keyspace-group", "finish-split", "0"} + args := []string{"-u", pdAddr, "keyspace-group", "finish-split", "1"} 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 { @@ -358,3 +354,95 @@ func TestMergeKeyspaceGroup(t *testing.T) { re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes")) re.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop")) } + +func TestKeyspaceGroupState(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) + for i := 0; i < 10; i++ { + keyspaces = append(keyspaces, fmt.Sprintf("keyspace_%d", i)) + } + tc, err := tests.NewTestAPICluster(ctx, 1, 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", "1"} + output, err := pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + args = []string{"-u", pdAddr, "keyspace-group", "--state", "split"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + var keyspaceGroups []*endpoint.KeyspaceGroup + err = json.Unmarshal(output, &keyspaceGroups) + re.NoError(err) + re.Len(keyspaceGroups, 0) + testutil.Eventually(re, func() bool { + args := []string{"-u", pdAddr, "keyspace-group", "split", "0", "2", "3"} + output, err := pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + return strings.Contains(string(output), "Success") + }) + args = []string{"-u", pdAddr, "keyspace-group", "--state", "split"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + err = json.Unmarshal(output, &keyspaceGroups) + re.NoError(err) + re.Len(keyspaceGroups, 2) + re.Equal(keyspaceGroups[0].ID, uint32(0)) + re.Equal(keyspaceGroups[1].ID, uint32(2)) + + args = []string{"-u", pdAddr, "keyspace-group", "finish-split", "2"} + 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", "--state", "merge"} + output, err = pdctl.ExecuteCommand(cmd, args...) + re.NoError(err) + strings.Contains(string(output), "Success") + err = json.Unmarshal(output, &keyspaceGroups) + re.NoError(err) + err = json.Unmarshal(output, &keyspaceGroups) + re.NoError(err) + re.Len(keyspaceGroups, 1) + re.Equal(keyspaceGroups[0].ID, uint32(0)) + + 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 b5ccaf01e2b..246bc2a60da 100644 --- a/tools/pd-ctl/pdctl/command/keyspace_group_command.go +++ b/tools/pd-ctl/pdctl/command/keyspace_group_command.go @@ -15,12 +15,15 @@ package command import ( + "encoding/json" "fmt" "net/http" "net/url" "strconv" + "strings" "github.com/spf13/cobra" + "github.com/tikv/pd/pkg/storage/endpoint" ) const keyspaceGroupsPrefix = "pd/api/v2/tso/keyspace-groups" @@ -28,9 +31,9 @@ const keyspaceGroupsPrefix = "pd/api/v2/tso/keyspace-groups" // NewKeyspaceGroupCommand return a keyspace group subcommand of rootCmd func NewKeyspaceGroupCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "keyspace-group ", - Short: "show keyspace group information with the given ID", - Run: showKeyspaceGroupCommandFunc, + Use: "keyspace-group [command] [flags]", + Short: "show keyspace group information", + Run: showKeyspaceGroupsCommandFunc, } cmd.AddCommand(newSplitKeyspaceGroupCommand()) cmd.AddCommand(newSplitRangeKeyspaceGroupCommand()) @@ -39,6 +42,7 @@ func NewKeyspaceGroupCommand() *cobra.Command { cmd.AddCommand(newFinishMergeKeyspaceGroupCommand()) cmd.AddCommand(newSetNodesKeyspaceGroupCommand()) cmd.AddCommand(newSetPriorityKeyspaceGroupCommand()) + cmd.Flags().String("state", "", "state filter") return cmd } @@ -107,16 +111,46 @@ func newSetPriorityKeyspaceGroupCommand() *cobra.Command { return r } -func showKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { - if len(args) < 1 { +func showKeyspaceGroupsCommandFunc(cmd *cobra.Command, args []string) { + prefix := keyspaceGroupsPrefix + if len(args) > 1 { cmd.Usage() return } - r, err := doRequest(cmd, fmt.Sprintf("%s/%s", keyspaceGroupsPrefix, args[0]), http.MethodGet, http.Header{}) + cFunc := convertToKeyspaceGroups + if len(args) == 1 { + if _, err := strconv.Atoi(args[0]); err != nil { + cmd.Println("keyspace_group_id should be a number") + return + } + prefix = fmt.Sprintf("%s/%s", keyspaceGroupsPrefix, args[0]) + cFunc = convertToKeyspaceGroup + } else { + flags := cmd.Flags() + state, err := flags.GetString("state") + if err != nil { + cmd.Printf("Failed to get state: %s\n", err) + } + stateValue := "" + state = strings.ToLower(state) + switch state { + case "merge", "split": + stateValue = fmt.Sprintf("state=%v", state) + default: + cmd.Println("Unknown state: " + state) + return + } + + if len(stateValue) != 0 { + prefix = fmt.Sprintf("%v?%v", keyspaceGroupsPrefix, stateValue) + } + } + r, err := doRequest(cmd, prefix, http.MethodGet, http.Header{}) if err != nil { cmd.Printf("Failed to get the keyspace groups information: %s\n", err) return } + r = cFunc(r) cmd.Println(r) } @@ -295,3 +329,29 @@ func setPriorityKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) { "Priority": priority, }) } + +func convertToKeyspaceGroup(content string) string { + kg := endpoint.KeyspaceGroup{} + err := json.Unmarshal([]byte(content), &kg) + if err != nil { + return content + } + byteArr, err := json.MarshalIndent(kg, "", " ") + if err != nil { + return content + } + return string(byteArr) +} + +func convertToKeyspaceGroups(content string) string { + kgs := []*endpoint.KeyspaceGroup{} + err := json.Unmarshal([]byte(content), &kgs) + if err != nil { + return content + } + byteArr, err := json.MarshalIndent(kgs, "", " ") + if err != nil { + return content + } + return string(byteArr) +}