support metrics HTTP interface for tso/resource manager server

Signed-off-by: Ryan Leung <[email protected]>
tikv · Apr 17, 2023 · 8b540af · 8b540af
1 parent 0e42a50
commit 8b540af
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 0 deletions.
diff --git a/pkg/mcs/resource_manager/server/apis/v1/api.go b/pkg/mcs/resource_manager/server/apis/v1/api.go
@@ -24,7 +24,9 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/joho/godotenv"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	rmserver "github.com/tikv/pd/pkg/mcs/resource_manager/server"
+	"github.com/tikv/pd/pkg/mcs/utils"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/apiutil/multiservicesapi"
 )
@@ -75,6 +77,7 @@ func NewService(srv *rmserver.Service) *Service {
 		c.Next()
 	})
 	apiHandlerEngine.Use(multiservicesapi.ServiceRedirector())
+	apiHandlerEngine.GET("metrics", utils.PromHandler(promhttp.Handler()))
 	endpoint := apiHandlerEngine.Group(APIPathPrefix)
 	s := &Service{
 		manager:          manager,

diff --git a/pkg/mcs/tso/server/apis/v1/api.go b/pkg/mcs/tso/server/apis/v1/api.go
@@ -22,7 +22,9 @@ import (
 	"github.com/gin-contrib/gzip"
 	"github.com/gin-gonic/gin"
 	"github.com/joho/godotenv"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	tsoserver "github.com/tikv/pd/pkg/mcs/tso/server"
+	"github.com/tikv/pd/pkg/mcs/utils"
 	"github.com/tikv/pd/pkg/tso"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/apiutil/multiservicesapi"
@@ -80,6 +82,7 @@ func NewService(srv *tsoserver.Service) *Service {
 		c.Next()
 	})
 	apiHandlerEngine.Use(multiservicesapi.ServiceRedirector())
+	apiHandlerEngine.GET("metrics", utils.PromHandler(promhttp.Handler()))
 	endpoint := apiHandlerEngine.Group(APIPathPrefix)
 	s := &Service{
 		srv:              srv,

diff --git a/pkg/mcs/utils/util.go b/pkg/mcs/utils/util.go
@@ -16,8 +16,10 @@ package utils
 
 import (
 	"context"
+	"net/http"
 	"time"
 
+	"github.com/gin-gonic/gin"
 	"github.com/pkg/errors"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"go.etcd.io/etcd/clientv3"
@@ -46,3 +48,10 @@ func InitClusterID(ctx context.Context, client *clientv3.Client) (id uint64, err
 	}
 	return 0, errors.Errorf("failed to init cluster ID after retrying %d times", maxRetryTimes)
 }
+
+// PromHandler is a handler to get prometheus metrics.
+func PromHandler(handler http.Handler) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		handler.ServeHTTP(c.Writer, c.Request)
+	}
+}
diff --git a/tests/integrations/mcs/resource_manager/server_test.go b/tests/integrations/mcs/resource_manager/server_test.go
@@ -15,6 +15,8 @@
 package resourcemanager_test
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"io"
@@ -92,4 +94,20 @@ func TestResourceManagerServer(t *testing.T) {
 		re.NoError(err)
 		re.Equal("{\"name\":\"pingcap\",\"mode\":1,\"r_u_settings\":{\"r_u\":{\"state\":{\"initialized\":false}}},\"priority\":0}", string(respString))
 	}
+
+	// Test metrics handler
+	{
+		resp, err := http.Get(addr + "/metrics")
+		re.NoError(err)
+		defer resp.Body.Close()
+		re.Equal(http.StatusOK, resp.StatusCode)
+		respString, err := io.ReadAll(resp.Body)
+		re.NoError(err)
+		reader := bytes.NewReader(respString)
+		gzipReader, err := gzip.NewReader(reader)
+		re.NoError(err)
+		output, err := io.ReadAll(gzipReader)
+		re.NoError(err)
+		re.Contains(string(output), "resource_manager_server_info")
+	}
 }
diff --git a/tests/integrations/mcs/tso/server_test.go b/tests/integrations/mcs/tso/server_test.go
@@ -15,8 +15,11 @@
 package tso
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"fmt"
+	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -345,3 +348,36 @@ func TestAdvertiseAddr(t *testing.T) {
 	tsoServerConf := s.GetConfig()
 	re.Equal(u, tsoServerConf.AdvertiseListenAddr)
 }
+
+func TestMetrics(t *testing.T) {
+	re := require.New(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := tests.NewTestAPICluster(ctx, 1)
+	defer cluster.Destroy()
+	re.NoError(err)
+
+	err = cluster.RunInitialServers()
+	re.NoError(err)
+
+	leaderName := cluster.WaitLeader()
+	leader := cluster.GetServer(leaderName)
+
+	u := tempurl.Alloc()
+	s, cleanup := mcs.StartSingleTSOTestServer(ctx, re, leader.GetAddr(), u)
+	defer cleanup()
+
+	resp, err := http.Get(s.GetConfig().GetAdvertiseListenAddr() + "/metrics")
+	re.NoError(err)
+	defer resp.Body.Close()
+	re.Equal(http.StatusOK, resp.StatusCode)
+	respString, err := io.ReadAll(resp.Body)
+	re.NoError(err)
+	reader := bytes.NewReader(respString)
+	gzipReader, err := gzip.NewReader(reader)
+	re.NoError(err)
+	output, err := io.ReadAll(gzipReader)
+	re.NoError(err)
+	re.Contains(string(output), "tso_server_info")
+}