Add API interface to obtain the TSO keyspace group member info

Signed-off-by: JmPotato <ghzpotato@gmail.com>

pkg/mcs/tso/server/apis/v1/api.go

@@ -22,16 +22,22 @@ import (
 	"github.com/gin-contrib/gzip"
 	"github.com/gin-gonic/gin"
 	"github.com/joho/godotenv"
+	"github.com/pingcap/kvproto/pkg/tsopb"
+	"github.com/pingcap/log"
 	tsoserver "github.com/tikv/pd/pkg/mcs/tso/server"
 	"github.com/tikv/pd/pkg/mcs/utils"
+	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/tso"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/apiutil/multiservicesapi"
 	"github.com/unrolled/render"
+	"go.uber.org/zap"
 )
 
-// APIPathPrefix is the prefix of the API path.
-const APIPathPrefix = "/tso/api/v1"
+const (
+	// APIPathPrefix is the prefix of the API path.
+	APIPathPrefix = "/tso/api/v1"
+)
 
 var (
 	once            sync.Once
@@ -46,14 +52,14 @@ var (
 func init() {
 	tsoserver.SetUpRestHandler = func(srv *tsoserver.Service) (http.Handler, apiutil.APIServiceGroup) {
 		s := NewService(srv)
-		return s.handler(), apiServiceGroup
+		return s.apiHandlerEngine, apiServiceGroup
 	}
 }
 
 // Service is the tso service.
 type Service struct {
 	apiHandlerEngine *gin.Engine
-	baseEndpoint     *gin.RouterGroup
+	root             *gin.RouterGroup
 
 	srv *tsoserver.Service
 	rd  *render.Render
@@ -77,30 +83,64 @@ func NewService(srv *tsoserver.Service) *Service {
 	apiHandlerEngine.Use(cors.Default())
 	apiHandlerEngine.Use(gzip.Gzip(gzip.DefaultCompression))
 	apiHandlerEngine.Use(func(c *gin.Context) {
-		c.Set("service", srv)
+		c.Set(multiservicesapi.ServiceContextKey, srv)
 		c.Next()
 	})
 	apiHandlerEngine.Use(multiservicesapi.ServiceRedirector())
 	apiHandlerEngine.GET("metrics", utils.PromHandler())
-	endpoint := apiHandlerEngine.Group(APIPathPrefix)
+	root := apiHandlerEngine.Group(APIPathPrefix)
 	s := &Service{
 		srv:              srv,
 		apiHandlerEngine: apiHandlerEngine,
-		baseEndpoint:     endpoint,
+		root:             root,
 		rd:               createIndentRender(),
 	}
-	s.RegisterRouter()
+	s.RegisterAdminRouter()
+	s.RegisterKeyspaceGroupRouter()
 	return s
 }
 
-// RegisterRouter registers the router of the service.
-func (s *Service) RegisterRouter() {
+// RegisterAdminRouter registers the router of the TSO admin handler.
+func (s *Service) RegisterAdminRouter() {
+	router := s.root.Group("admin")
 	tsoAdminHandler := tso.NewAdminHandler(s.srv.GetHandler(), s.rd)
-	s.baseEndpoint.POST("/admin/reset-ts", gin.WrapF(tsoAdminHandler.ResetTS))
+	router.POST("/admin/reset-ts", gin.WrapF(tsoAdminHandler.ResetTS))
 }
 
-func (s *Service) handler() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		s.apiHandlerEngine.ServeHTTP(w, r)
-	})
+// RegisterKeyspaceGroupRouter registers the router of the TSO keyspace group handler.
+func (s *Service) RegisterKeyspaceGroupRouter() {
+	router := s.root.Group("keyspace-groups")
+	router.GET("/members", GetKeyspaceGroupMembers)
+}
+
+// KeyspaceGroupMember contains the keyspace group and its member information.
+type KeyspaceGroupMember struct {
+	Group     *endpoint.KeyspaceGroup
+	Member    *tsopb.Participant
+	IsPrimary bool   `json:"is_primary"`
+	PrimaryID uint64 `json:"primary_id"`
+}
+
+// GetKeyspaceGroupMembers gets the keyspace group members that the TSO service is serving.
+func GetKeyspaceGroupMembers(c *gin.Context) {
+	svr := c.MustGet(multiservicesapi.ServiceContextKey).(*tsoserver.Service)
+	kgm := svr.GetKeyspaceGroupManager()
+	keyspaceGroups := kgm.GetKeyspaceGroups()
+	members := make(map[uint32]*KeyspaceGroupMember, len(keyspaceGroups))
+	for id, group := range keyspaceGroups {
+		am, err := kgm.GetAllocatorManager(id)
+		if err != nil {
+			log.Error("failed to get allocator manager",
+				zap.Uint32("keyspace-group-id", id), zap.Error(err))
+			continue
+		}
+		member := am.GetMember()
+		members[id] = &KeyspaceGroupMember{
+			Group:     group,
+			Member:    member.GetMember().(*tsopb.Participant),
+			IsPrimary: member.IsLeader(),
+			PrimaryID: member.GetLeaderID(),
+		}
+	}
+	c.IndentedJSON(http.StatusOK, members)
 }

pkg/mcs/tso/server/server.go

@@ -273,6 +273,11 @@ func (s *Server) IsClosed() bool {
 	return atomic.LoadInt64(&s.isRunning) == 0
 }
 
+// GetKeyspaceGroupManager returns the manager of keyspace group.
+func (s *Server) GetKeyspaceGroupManager() *tso.KeyspaceGroupManager {
+	return s.keyspaceGroupManager
+}
+
 // GetTSOAllocatorManager returns the manager of TSO Allocator.
 func (s *Server) GetTSOAllocatorManager(keyspaceGroupID uint32) (*tso.AllocatorManager, error) {
 	return s.keyspaceGroupManager.GetAllocatorManager(keyspaceGroupID)

pkg/tso/keyspace_group_manager.go

@@ -826,6 +826,20 @@ func (kgm *KeyspaceGroupManager) GetElectionMember(
 	return am.GetMember(), nil
 }
 
+// GetKeyspaceGroups returns all keyspace groups managed by the current keyspace group manager.
+func (kgm *KeyspaceGroupManager) GetKeyspaceGroups() map[uint32]*endpoint.KeyspaceGroup {
+	kgm.RLock()
+	defer kgm.RUnlock()
+	keyspaceGroups := make(map[uint32]*endpoint.KeyspaceGroup)
+	for _, keyspaceGroupID := range kgm.keyspaceLookupTable {
+		if _, ok := keyspaceGroups[keyspaceGroupID]; ok {
+			continue
+		}
+		keyspaceGroups[keyspaceGroupID] = kgm.kgs[keyspaceGroupID]
+	}
+	return keyspaceGroups
+}
+
 // HandleTSORequest forwards TSO allocation requests to correct TSO Allocators of the given keyspace group.
 func (kgm *KeyspaceGroupManager) HandleTSORequest(
 	keyspaceID, keyspaceGroupID uint32,

pkg/utils/apiutil/multiservicesapi/middleware.go

@@ -26,16 +26,19 @@ import (
 	"go.uber.org/zap"
 )
 
-// HTTP headers.
 const (
+	// ServiceAllowDirectHandle is the header key to allow direct handle.
 	ServiceAllowDirectHandle = "service-allow-direct-handle"
-	ServiceRedirectorHeader  = "service-redirector"
+	// ServiceRedirectorHeader is the header key to indicate the request is redirected.
+	ServiceRedirectorHeader = "service-redirector"
+	// ServiceContextKey is the key to get service server from gin.Context.
+	ServiceContextKey = "service"
 )
 
 // ServiceRedirector is a middleware to redirect the request to the right place.
 func ServiceRedirector() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		svr := c.MustGet("service").(bs.Server)
+		svr := c.MustGet(ServiceContextKey).(bs.Server)
 		allowDirectHandle := len(c.Request.Header.Get(ServiceAllowDirectHandle)) > 0
 		isServing := svr.IsServing()
 		if allowDirectHandle || isServing {

tests/integrations/mcs/tso/api_test.go

@@ -0,0 +1,106 @@
+// 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 tso
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+	tso "github.com/tikv/pd/pkg/mcs/tso/server"
+	apis "github.com/tikv/pd/pkg/mcs/tso/server/apis/v1"
+	mcsutils "github.com/tikv/pd/pkg/mcs/utils"
+	"github.com/tikv/pd/tests"
+	"github.com/tikv/pd/tests/integrations/mcs"
+)
+
+const (
+	tsoKeyspaceGroupsPrefix = "/tso/api/v1/keyspace-groups"
+)
+
+// dialClient used to dial http request.
+var dialClient = &http.Client{
+	Transport: &http.Transport{
+		DisableKeepAlives: true,
+	},
+}
+
+type tsoAPITestSuite struct {
+	suite.Suite
+	ctx        context.Context
+	cancel     context.CancelFunc
+	pdCluster  *tests.TestCluster
+	tsoCluster *mcs.TestTSOCluster
+}
+
+func TestTSOAPI(t *testing.T) {
+	suite.Run(t, new(tsoAPITestSuite))
+}
+
+func (suite *tsoAPITestSuite) SetupTest() {
+	re := suite.Require()
+
+	var err error
+	suite.ctx, suite.cancel = context.WithCancel(context.Background())
+	suite.pdCluster, err = tests.NewTestAPICluster(suite.ctx, 1)
+	re.NoError(err)
+	err = suite.pdCluster.RunInitialServers()
+	re.NoError(err)
+	leaderName := suite.pdCluster.WaitLeader()
+	pdLeaderServer := suite.pdCluster.GetServer(leaderName)
+	re.NoError(pdLeaderServer.BootstrapCluster())
+	suite.tsoCluster, err = mcs.NewTestTSOCluster(suite.ctx, 1, pdLeaderServer.GetAddr())
+	re.NoError(err)
+}
+
+func (suite *tsoAPITestSuite) TearDownTest() {
+	suite.cancel()
+	suite.tsoCluster.Destroy()
+	suite.pdCluster.Destroy()
+}
+
+func (suite *tsoAPITestSuite) TestGetKeyspaceGroupMembers() {
+	re := suite.Require()
+
+	primary := suite.tsoCluster.WaitForDefaultPrimaryServing(re)
+	re.NotNil(primary)
+	members := mustGetKeyspaceGroupMembers(re, primary)
+	re.Len(members, 1)
+	defaultGroupMember := members[mcsutils.DefaultKeyspaceGroupID]
+	re.NotNil(defaultGroupMember)
+	re.Equal(mcsutils.DefaultKeyspaceGroupID, defaultGroupMember.Group.ID)
+	re.True(defaultGroupMember.IsPrimary)
+	primaryMember, err := primary.GetMember(mcsutils.DefaultKeyspaceID, mcsutils.DefaultKeyspaceGroupID)
+	re.NoError(err)
+	re.Equal(primaryMember.GetLeaderID(), defaultGroupMember.PrimaryID)
+}
+
+func mustGetKeyspaceGroupMembers(re *require.Assertions, server *tso.Server) map[uint32]*apis.KeyspaceGroupMember {
+	httpReq, err := http.NewRequest(http.MethodGet, server.GetAddr()+tsoKeyspaceGroupsPrefix+"/members", nil)
+	re.NoError(err)
+	httpResp, err := dialClient.Do(httpReq)
+	re.NoError(err)
+	defer httpResp.Body.Close()
+	data, err := io.ReadAll(httpResp.Body)
+	re.NoError(err)
+	re.Equal(http.StatusOK, httpResp.StatusCode, string(data))
+	var resp map[uint32]*apis.KeyspaceGroupMember
+	re.NoError(json.Unmarshal(data, &resp))
+	return resp
+}