pdctl: support show keyspace group primary

Signed-off-by: lhy1024 <[email protected]>

pkg/keyspace/tso_keyspace_group.go

@@ -17,6 +17,7 @@ package keyspace
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"sort"
 	"strconv"
 	"strings"
@@ -25,6 +26,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/tsopb"
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/balancer"
 	"github.com/tikv/pd/pkg/mcs/discovery"
@@ -1010,3 +1012,23 @@ func (m *GroupManager) FinishMergeKeyspaceByID(mergeTargetID uint32) error {
 		zap.Reflect("merge-list", mergeList))
 	return nil
 }
+
+func (m *GroupManager) GetKeyspaceGroupPrimaryByID(id uint32) (string, error) {
+	// default keyspace group: "/ms/{cluster_id}/tso/00000/primary".
+	// non-default keyspace group: "/ms/{cluster_id}/tso/keyspace_groups/election/{group}/primary".
+	path := fmt.Sprintf("/ms/%d/tso/00000/primary", m.clusterID)
+	if id != utils.DefaultKeyspaceGroupID {
+		path = fmt.Sprintf("/ms/%d/tso/keyspace_groups/election/%05d/primary", m.clusterID, id)
+	}
+	leader := &tsopb.Participant{}
+	ok, _, err := etcdutil.GetProtoMsgWithModRev(m.client, path, leader)
+	if err != nil {
+		return "", err
+	}
+	if !ok {
+		return "", ErrKeyspaceGroupPrimaryNotFound
+	}
+	// The format of leader name is address-groupID.
+	contents := strings.Split(leader.GetName(), "-")
+	return contents[0], err
+}

pkg/keyspace/util.go

@@ -94,6 +94,9 @@ var (
 	}
 	// Only keyspaces in the state specified by allowChangeConfig are allowed to change their config.
 	allowChangeConfig = []keyspacepb.KeyspaceState{keyspacepb.KeyspaceState_ENABLED, keyspacepb.KeyspaceState_DISABLED}
+
+	// ErrKeyspaceGroupPrimaryNotFound is used to indicate primary of target keyspace group does not exist.
+	ErrKeyspaceGroupPrimaryNotFound = errors.New("primary of keyspace group does not exist")
 )
 
 // validateID check if keyspace falls within the acceptable range.

pkg/tso/keyspace_group_manager.go

@@ -683,6 +683,7 @@ func (kgm *KeyspaceGroupManager) updateKeyspaceGroup(group *endpoint.KeyspaceGro
 	}
 
 	// If the keyspace group is not initialized, initialize it.
+	// The format of leader name is address-groupID.
 	uniqueName := fmt.Sprintf("%s-%05d", kgm.electionNamePrefix, group.ID)
 	uniqueID := memberutil.GenerateUniqueID(uniqueName)
 	log.Info("joining primary election",

server/apiv2/handlers/tso_keyspace_group.go

@@ -47,6 +47,7 @@ func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) {
 	router.DELETE("/:id/split", FinishSplitKeyspaceByID)
 	router.POST("/:id/merge", MergeKeyspaceGroups)
 	router.DELETE("/:id/merge", FinishMergeKeyspaceByID)
+	router.GET("/:id/primary", GetKeyspaceGroupPrimaryByID)
 }
 
 // CreateKeyspaceGroupParams defines the params for creating keyspace groups.
@@ -485,6 +486,33 @@ func SetPriorityForKeyspaceGroup(c *gin.Context) {
 	c.JSON(http.StatusOK, nil)
 }
 
+func GetKeyspaceGroupPrimaryByID(c *gin.Context) {
+	id, err := validateKeyspaceGroupID(c)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id")
+		return
+	}
+	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
+	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, groupManagerUninitializedErr)
+		return
+	}
+	// check if keyspace group exists
+	kg, err := manager.GetKeyspaceGroupByID(id)
+	if err != nil || kg == nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, "keyspace group does not exist")
+		return
+	}
+	// get primary
+	primary, err := manager.GetKeyspaceGroupPrimaryByID(id)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+		return
+	}
+	c.JSON(http.StatusOK, primary)
+}
+
 func validateKeyspaceGroupID(c *gin.Context) (uint32, error) {
 	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
 	if err != nil {

tests/pdctl/keyspace/keyspace_group_test.go

@@ -450,3 +450,98 @@ func TestKeyspaceGroupState(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 TestShowKeyspaceGroupPrimary(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/pkg/tso/fastGroupSplitPatroller", `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, 3, func(conf *config.Config, serverName string) {
+		conf.Keyspace.PreAlloc = keyspaces
+	})
+	re.NoError(err)
+	err = tc.RunInitialServers()
+	re.NoError(err)
+	pdAddr := tc.GetConfig().GetClientURL()
+
+	s1, tsoServerCleanup1, err := tests.StartSingleTSOTestServer(ctx, re, pdAddr, tempurl.Alloc())
+	defer tsoServerCleanup1()
+	re.NoError(err)
+	s2, 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())
+	defaultKeyspaceGroupID := fmt.Sprintf("%d", utils.DefaultKeyspaceGroupID)
+
+	// check keyspace group 0 information.
+	var keyspaceGroup endpoint.KeyspaceGroup
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace-group"}
+		output, err := pdctl.ExecuteCommand(cmd, append(args, defaultKeyspaceGroupID)...)
+		re.NoError(err)
+
+		err = json.Unmarshal(output, &keyspaceGroup)
+		re.NoError(err)
+		re.Equal(utils.DefaultKeyspaceGroupID, keyspaceGroup.ID)
+		return len(keyspaceGroup.Members) == 2
+	})
+	for _, member := range keyspaceGroup.Members {
+		re.Contains([]string{s1.GetAddr(), s2.GetAddr()}, member.Address)
+	}
+
+	// get primary for keyspace group 0.
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace-group", "primary", defaultKeyspaceGroupID}
+		output, err := pdctl.ExecuteCommand(cmd, args...)
+		re.NoError(err)
+		addr := strings.ReplaceAll(string(output), "\"", "")
+		addr = strings.ReplaceAll(addr, "\n", "")
+		return s1.GetAddr() == addr || s2.GetAddr() == addr
+	})
+
+	// 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")
+	})
+
+	// check keyspace group 1 information.
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace-group"}
+		output, err := pdctl.ExecuteCommand(cmd, append(args, "1")...)
+		re.NoError(err)
+
+		err = json.Unmarshal(output, &keyspaceGroup)
+		re.NoError(err)
+		return len(keyspaceGroup.Members) == 2
+	})
+	for _, member := range keyspaceGroup.Members {
+		re.Contains([]string{s1.GetAddr(), s2.GetAddr()}, member.Address)
+	}
+
+	// get primary for keyspace group 1.
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace-group", "primary", "1"}
+		output, err := pdctl.ExecuteCommand(cmd, args...)
+		re.NoError(err)
+		addr := strings.ReplaceAll(string(output), "\"", "")
+		addr = strings.ReplaceAll(addr, "\n", "")
+		return s1.GetAddr() == addr || s2.GetAddr() == addr
+	})
+
+	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes"))
+	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/fastGroupSplitPatroller"))
+	re.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop"))
+}

tools/pd-ctl/pdctl/command/keyspace_group_command.go

@@ -42,6 +42,7 @@ func NewKeyspaceGroupCommand() *cobra.Command {
 	cmd.AddCommand(newFinishMergeKeyspaceGroupCommand())
 	cmd.AddCommand(newSetNodesKeyspaceGroupCommand())
 	cmd.AddCommand(newSetPriorityKeyspaceGroupCommand())
+	cmd.AddCommand(newShowKeyspaceGroupPrimaryCommand())
 	cmd.Flags().String("state", "", "state filter")
 	return cmd
 }
@@ -111,6 +112,15 @@ func newSetPriorityKeyspaceGroupCommand() *cobra.Command {
 	return r
 }
 
+func newShowKeyspaceGroupPrimaryCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "primary <keyspace_group_id>",
+		Short: "show th primary of tso nodes for keyspace group with the given ID.",
+		Run:   showKeyspaceGroupPrimaryCommandFunc,
+	}
+	return r
+}
+
 func showKeyspaceGroupsCommandFunc(cmd *cobra.Command, args []string) {
 	prefix := keyspaceGroupsPrefix
 	if len(args) > 1 {
@@ -337,6 +347,24 @@ func setPriorityKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) {
 	})
 }
 
+func showKeyspaceGroupPrimaryCommandFunc(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
+	}
+	r, err := doRequest(cmd, fmt.Sprintf("%s/%s/primary", keyspaceGroupsPrefix, args[0]), http.MethodGet, http.Header{})
+	if err != nil {
+		cmd.Printf("Failed to get the keyspace group primary information: %s\n", err)
+		return
+	}
+	cmd.Println(r)
+}
+
 func convertToKeyspaceGroup(content string) string {
 	kg := endpoint.KeyspaceGroup{}
 	err := json.Unmarshal([]byte(content), &kg)