Skip to content

Commit

Permalink
tools: support querying keyspace group by state (tikv#6706)
Browse files Browse the repository at this point in the history 
ref tikv#5895

Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
  • Loading branch information
@rleungx @ti-chi-bot
rleungx and ti-chi-bot[bot] committed Aug 2, 2023
1 parent d1abae5 commit 7c6cc4f
Show file tree
Hide file tree
Showing 3 changed files with 184 additions and 13 deletions.
25 changes: 24 additions & 1 deletion server/apiv2/handlers/tso_keyspace_group.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ package handlers
import (
"net/http"
"strconv"
"strings"
"sync"

"github.com/gin-gonic/gin"
Expand Down Expand Up @@ -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.
Expand Down
100 changes: 94 additions & 6 deletions tests/pdctl/keyspace/keyspace_group_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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"))
}
72 changes: 66 additions & 6 deletions tools/pd-ctl/pdctl/command/keyspace_group_command.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,22 +15,25 @@
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"

// NewKeyspaceGroupCommand return a keyspace group subcommand of rootCmd
func NewKeyspaceGroupCommand() *cobra.Command {
cmd := &cobra.Command{
Use: "keyspace-group <keyspace_group_id>",
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())
Expand All @@ -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
}

Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -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)
}

0 comments on commit 7c6cc4f

Please sign in to comment.