pdctl: support show keyspace meta with refresh group id (#6751)

close #6746

Signed-off-by: lhy1024 <admin@liudos.us>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

pkg/keyspace/keyspace.go

@@ -48,6 +48,8 @@ const (
 	// UserKindKey is the key for user kind in keyspace config.
 	UserKindKey = "user_kind"
 	// TSOKeyspaceGroupIDKey is the key for tso keyspace group id in keyspace config.
+	// Note: Config[TSOKeyspaceGroupIDKey] is only used to judge whether there is keyspace group id.
+	// It will not update the keyspace group id when merging or splitting.
 	TSOKeyspaceGroupIDKey = "tso_keyspace_group_id"
 	// MaxEtcdTxnOps is the max value of operations in an etcd txn. The default limit of etcd txn op is 128.
 	// We use 120 here to leave some space for other operations.

pkg/keyspace/tso_keyspace_group.go

@@ -388,6 +388,20 @@ func (m *GroupManager) getKeyspaceConfigByKindLocked(userKind endpoint.UserKind)
 	return config, nil
 }
 
+// GetGroupByKeyspaceID returns the keyspace group ID for the given keyspace ID.
+func (m *GroupManager) GetGroupByKeyspaceID(id uint32) (uint32, error) {
+	m.RLock()
+	defer m.RUnlock()
+	for _, groups := range m.groups {
+		for _, group := range groups.GetAll() {
+			if slice.Contains(group.Keyspaces, id) {
+				return group.ID, nil
+			}
+		}
+	}
+	return 0, ErrKeyspaceNotInAnyKeyspaceGroup
+}
+
 var failpointOnce sync.Once
 
 // UpdateKeyspaceForGroup updates the keyspace field for the keyspace group.

pkg/keyspace/util.go

@@ -70,6 +70,8 @@ var (
 	}
 	// ErrKeyspaceNotInKeyspaceGroup is used to indicate target keyspace is not in this keyspace group.
 	ErrKeyspaceNotInKeyspaceGroup = errors.New("keyspace is not in this keyspace group")
+	// ErrKeyspaceNotInAnyKeyspaceGroup is used to indicate target keyspace is not in any keyspace group.
+	ErrKeyspaceNotInAnyKeyspaceGroup = errors.New("keyspace is not in any keyspace group")
 	// ErrNodeNotInKeyspaceGroup is used to indicate the tso node is not in this keyspace group.
 	ErrNodeNotInKeyspaceGroup = errors.New("the tso node is not in this keyspace group")
 	// ErrKeyspaceGroupNotEnoughReplicas is used to indicate not enough replicas in the keyspace group.

server/apiv2/handlers/keyspace.go

@@ -110,6 +110,20 @@ func LoadKeyspace(c *gin.Context) {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
 		return
 	}
+	if value, ok := c.GetQuery("force_refresh_group_id"); ok && value == "true" {
+		groupManager := svr.GetKeyspaceGroupManager()
+		if groupManager == nil {
+			c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr)
+			return
+		}
+		// keyspace has been checked in LoadKeyspace, so no need to check again.
+		groupID, err := groupManager.GetGroupByKeyspaceID(meta.GetId())
+		if err != nil {
+			c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+			return
+		}
+		meta.Config[keyspace.TSOKeyspaceGroupIDKey] = strconv.FormatUint(uint64(groupID), 10)
+	}
 	c.IndentedJSON(http.StatusOK, &KeyspaceMeta{meta})
 }
 

tests/pdctl/keyspace/keyspace_group_test.go

@@ -561,9 +561,8 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 		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)
+		re.NoErrorf(err, "output: %s", string(output))
 		return len(keyspaceGroup.Members) == 2
 	})
 	for _, member := range keyspaceGroup.Members {

tests/pdctl/keyspace/keyspace_test.go

@@ -0,0 +1,103 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package keyspace_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/pingcap/failpoint"
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/keyspace"
+	"github.com/tikv/pd/pkg/mcs/utils"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	api "github.com/tikv/pd/server/apiv2/handlers"
+	"github.com/tikv/pd/server/config"
+	"github.com/tikv/pd/tests"
+	"github.com/tikv/pd/tests/pdctl"
+	pdctlCmd "github.com/tikv/pd/tools/pd-ctl/pdctl"
+)
+
+func TestKeyspace(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 := 1; 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()
+
+	ttc, err := tests.NewTestTSOCluster(ctx, 2, pdAddr)
+	re.NoError(err)
+	defer ttc.Destroy()
+	cmd := pdctlCmd.GetRootCmd()
+
+	tc.WaitLeader()
+	leaderServer := tc.GetServer(tc.GetLeader())
+	re.NoError(leaderServer.BootstrapCluster())
+	defaultKeyspaceGroupID := fmt.Sprintf("%d", utils.DefaultKeyspaceGroupID)
+
+	var k api.KeyspaceMeta
+	keyspaceName := "keyspace_1"
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace", keyspaceName}
+		output, err := pdctl.ExecuteCommand(cmd, args...)
+		re.NoError(err)
+		re.NoError(json.Unmarshal(output, &k))
+		return k.GetName() == keyspaceName
+	})
+	re.Equal(uint32(1), k.GetId())
+	re.Equal(defaultKeyspaceGroupID, k.Config[keyspace.TSOKeyspaceGroupIDKey])
+
+	// split keyspace group.
+	newGroupID := "2"
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace-group", "split", "0", newGroupID, "1"}
+		output, err := pdctl.ExecuteCommand(cmd, args...)
+		re.NoError(err)
+		return strings.Contains(string(output), "Success")
+	})
+
+	// check keyspace group in keyspace whether changed.
+	testutil.Eventually(re, func() bool {
+		args := []string{"-u", pdAddr, "keyspace", keyspaceName}
+		output, err := pdctl.ExecuteCommand(cmd, args...)
+		re.NoError(err)
+		re.NoError(json.Unmarshal(output, &k))
+		return newGroupID == k.Config[keyspace.TSOKeyspaceGroupIDKey]
+	})
+
+	// test error name
+	args := []string{"-u", pdAddr, "keyspace", "error_name"}
+	output, err := pdctl.ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	re.Contains(string(output), "Fail")
+	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_command.go

@@ -0,0 +1,48 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package command
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/spf13/cobra"
+)
+
+const keyspacePrefix = "pd/api/v2/keyspaces"
+
+// NewKeyspaceCommand returns a keyspace subcommand of rootCmd.
+func NewKeyspaceCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "keyspace [command] [flags]",
+		Short: "show keyspace information",
+		Run:   showKeyspaceCommandFunc,
+	}
+	return cmd
+}
+
+func showKeyspaceCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) != 1 {
+		cmd.Usage()
+		return
+	}
+
+	resp, err := doRequest(cmd, fmt.Sprintf("%s/%s?force_refresh_group_id=true", keyspacePrefix, args[0]), http.MethodGet, http.Header{})
+	if err != nil {
+		cmd.Printf("Failed to get the keyspace information: %s\n", err)
+		return
+	}
+	cmd.Println(resp)
+}

tools/pd-ctl/pdctl/ctl.go

@@ -66,6 +66,7 @@ func GetRootCmd() *cobra.Command {
 		command.NewCompletionCommand(),
 		command.NewUnsafeCommand(),
 		command.NewKeyspaceGroupCommand(),
+		command.NewKeyspaceCommand(),
 	)
 
 	rootCmd.Flags().ParseErrorsWhitelist.UnknownFlags = true