From 59dde6ac423ec9afd948245656b607582a94fd79 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Thu, 8 Feb 2024 12:33:33 +0800
Subject: [PATCH] add primary api for micro service

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 client/http/api.go                            |  5 ++++
 client/http/interface.go                      | 12 ++++++++++
 client/http/request_info.go                   |  1 +
 server/apiv2/handlers/micro_service.go        | 23 +++++++++++++++++++
 tests/integrations/mcs/members/member_test.go | 11 +++++++++
 5 files changed, 52 insertions(+)

diff --git a/client/http/api.go b/client/http/api.go
index 2cd7fee89c2..398307ec1ba 100644
--- a/client/http/api.go
+++ b/client/http/api.go
@@ -196,3 +196,8 @@ func PProfGoroutineWithDebugLevel(level int) string {
 func MicroServiceMembers(service string) string {
 	return fmt.Sprintf("%s/members/%s", microServicePrefix, service)
 }
+
+// MicroServicePrimary returns the path of PD HTTP API to get the primary of microservice.
+func MicroServicePrimary(service string) string {
+	return fmt.Sprintf("%s/primary/%s", microServicePrefix, service)
+}
diff --git a/client/http/interface.go b/client/http/interface.go
index 47bf5cc58e1..c7bfc37958d 100644
--- a/client/http/interface.go
+++ b/client/http/interface.go
@@ -92,6 +92,7 @@ type Client interface {
 	GetPDVersion(context.Context) (string, error)
 	/* Micro Service interfaces */
 	GetMicroServiceMembers(context.Context, string) ([]string, error)
+	GetMicroServicePrimary(context.Context, string) (string, error)
 	DeleteOperators(context.Context) error
 
 	/* Client-related methods */
@@ -868,6 +869,17 @@ func (c *client) GetMicroServiceMembers(ctx context.Context, service string) ([]
 	return members, nil
 }
 
+// GetMicroServicePrimary gets the primary of the microservice.
+func (c *client) GetMicroServicePrimary(ctx context.Context, service string) (string, error) {
+	var primary string
+	err := c.request(ctx, newRequestInfo().
+		WithName(getMicroServicePrimaryName).
+		WithURI(MicroServicePrimary(service)).
+		WithMethod(http.MethodGet).
+		WithResp(&primary))
+	return primary, err
+}
+
 // GetPDVersion gets the release version of the PD binary.
 func (c *client) GetPDVersion(ctx context.Context) (string, error) {
 	var ver struct {
diff --git a/client/http/request_info.go b/client/http/request_info.go
index e5e8f57724c..a7a6e9b7652 100644
--- a/client/http/request_info.go
+++ b/client/http/request_info.go
@@ -70,6 +70,7 @@ const (
 	accelerateScheduleInBatchName           = "AccelerateScheduleInBatch"
 	getMinResolvedTSByStoresIDsName         = "GetMinResolvedTSByStoresIDs"
 	getMicroServiceMembersName              = "GetMicroServiceMembers"
+	getMicroServicePrimaryName              = "GetMicroServicePrimary"
 	getPDVersionName                        = "GetPDVersion"
 	resetTSName                             = "ResetTS"
 	resetBaseAllocIDName                    = "ResetBaseAllocID"
diff --git a/server/apiv2/handlers/micro_service.go b/server/apiv2/handlers/micro_service.go
index 3c2be3748d4..48310d26dea 100644
--- a/server/apiv2/handlers/micro_service.go
+++ b/server/apiv2/handlers/micro_service.go
@@ -28,6 +28,7 @@ func RegisterMicroService(r *gin.RouterGroup) {
 	router := r.Group("ms")
 	router.Use(middlewares.BootstrapChecker())
 	router.GET("members/:service", GetMembers)
+	router.GET("primary/:service", GetPrimary)
 }
 
 // GetMembers gets all members of the cluster for the specified service.
@@ -55,3 +56,25 @@ func GetMembers(c *gin.Context) {
 
 	c.AbortWithStatusJSON(http.StatusInternalServerError, "please specify service")
 }
+
+// GetPrimary gets the primary member of the specified service.
+// @Tags     primary
+// @Summary  Get the primary member of the specified service.
+// @Produce  json
+// @Success  200  {object}  string
+// @Router   /ms/primary/{service} [get]
+func GetPrimary(c *gin.Context) {
+	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
+	if !svr.IsAPIServiceMode() {
+		c.AbortWithStatusJSON(http.StatusServiceUnavailable, "not support micro service")
+		return
+	}
+
+	if service := c.Param("service"); len(service) > 0 {
+		addr, _ := svr.GetServicePrimaryAddr(c.Request.Context(), service)
+		c.IndentedJSON(http.StatusOK, addr)
+		return
+	}
+
+	c.AbortWithStatusJSON(http.StatusInternalServerError, "please specify service")
+}
diff --git a/tests/integrations/mcs/members/member_test.go b/tests/integrations/mcs/members/member_test.go
index 41cd30bb5b7..b3d2a8048de 100644
--- a/tests/integrations/mcs/members/member_test.go
+++ b/tests/integrations/mcs/members/member_test.go
@@ -100,3 +100,14 @@ func (suite *memberTestSuite) TestMembers() {
 	re.NoError(err)
 	re.Len(members, 3)
 }
+
+func (suite *memberTestSuite) TestPrimary() {
+	re := suite.Require()
+	primary, err := suite.dialClient.GetMicroServicePrimary(suite.ctx, "tso")
+	re.NoError(err)
+	re.NotEmpty(primary)
+
+	primary, err = suite.dialClient.GetMicroServicePrimary(suite.ctx, "scheduling")
+	re.NoError(err)
+	re.NotEmpty(primary)
+}