add keyspace, we need Config[TSOKeyspaceGroupIDKey]?

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

pkg/keyspace/tso_keyspace_group.go

@@ -387,6 +387,19 @@ func (m *GroupManager) getKeyspaceConfigByKindLocked(userKind endpoint.UserKind)
 	return config, nil
 }
 
+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, errors.Errorf("keyspace %d not found", id)
+}
+
 var failpointOnce sync.Once
 
 // UpdateKeyspaceForGroup updates the keyspace field for the keyspace group.

server/apiv2/handlers/keyspace.go

@@ -42,6 +42,7 @@ func RegisterKeyspace(r *gin.RouterGroup) {
 	router.PATCH("/:name/config", UpdateKeyspaceConfig)
 	router.PUT("/:name/state", UpdateKeyspaceState)
 	router.GET("/id/:id", LoadKeyspaceByID)
+	router.GET("/:name/group", GetKeyspaceGroupID)
 }
 
 // CreateKeyspaceParams represents parameters needed when creating a new keyspace.
@@ -402,3 +403,30 @@ func (meta *KeyspaceMeta) UnmarshalJSON(data []byte) error {
 	meta.KeyspaceMeta = pbMeta
 	return nil
 }
+
+// GetKeyspaceGroupID returns the group that contains the given keyspace id.
+func GetKeyspaceGroupID(c *gin.Context) {
+	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
+	groupManager := svr.GetKeyspaceGroupManager()
+	if groupManager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr)
+		return
+	}
+	keyspaceManager := svr.GetKeyspaceManager()
+	if keyspaceManager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr)
+		return
+	}
+	name := c.Param("name")
+	meta, err := keyspaceManager.LoadKeyspace(name)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+		return
+	}
+	group, err := groupManager.GetGroupByKeyspaceID(meta.GetId())
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+		return
+	}
+	c.String(http.StatusOK, strconv.FormatUint(uint64(group), 10))
+}

tests/pdctl/keyspace/keyspace_test.go

@@ -0,0 +1,127 @@
+// 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"
+	"time"
+
+	"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/storage/endpoint"
+	"github.com/tikv/pd/pkg/utils/tempurl"
+	"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()
+
+	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)
+	}
+
+	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"
+	// split keyspace group.
+	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")
+	})
+	time.Sleep(30 * time.Second)
+
+	// check keyspace group in keyspace.
+	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 return 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", 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