Improve HTTP server include metrics and RESTful API (#28045)

/kind improvement
issue: #27653

Signed-off-by: PowderLi <min.li@zilliz.com>

internal/http/server.go

@@ -36,6 +36,11 @@ const (
 	ListenPortEnvKey  = "METRICS_PORT"
 )
 
+var (
+	metricsServer *http.ServeMux
+	server        *http.Server
+)
+
 type Handler struct {
 	Path        string
 	HandlerFunc http.HandlerFunc
@@ -53,7 +58,6 @@ func registerDefaults() {
 		Path:    HealthzRouterPath,
 		Handler: healthz.Handler(),
 	})
-
 	Register(&Handler{
 		Path:    EventLogRouterPath,
 		Handler: eventlog.Handler(),
@@ -61,12 +65,19 @@ func registerDefaults() {
 }
 
 func Register(h *Handler) {
+	if metricsServer == nil {
+		if paramtable.Get().HTTPCfg.EnablePprof.GetAsBool() {
+			metricsServer = http.DefaultServeMux
+		} else {
+			metricsServer = http.NewServeMux()
+		}
+	}
 	if h.HandlerFunc != nil {
-		http.HandleFunc(h.Path, h.HandlerFunc)
+		metricsServer.HandleFunc(h.Path, h.HandlerFunc)
 		return
 	}
 	if h.Handler != nil {
-		http.Handle(h.Path, h.Handler)
+		metricsServer.Handle(h.Path, h.Handler)
 	}
 }
 
@@ -75,7 +86,7 @@ func ServeHTTP() {
 	go func() {
 		bindAddr := getHTTPAddr()
 		log.Info("management listen", zap.String("addr", bindAddr))
-		server := &http.Server{Addr: bindAddr, ReadTimeout: 10 * time.Second}
+		server = &http.Server{Handler: metricsServer, Addr: bindAddr, ReadTimeout: 10 * time.Second}
 		if err := server.ListenAndServe(); err != nil {
 			log.Error("handle metrics failed", zap.Error(err))
 		}

internal/http/server_test.go

@@ -20,12 +20,13 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
-	"net/http/httptest"
+	"os"
+	"strings"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/zap"
 
@@ -35,29 +36,25 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
 
-func TestMain(t *testing.M) {
-	paramtable.Init()
-}
-
-func TestGetHTTPAddr(t *testing.T) {
-	assert.Equal(t, getHTTPAddr(), ":"+DefaultListenPort)
-	testPort := "9092"
-	t.Setenv(ListenPortEnvKey, testPort)
-	assert.Equal(t, getHTTPAddr(), ":"+testPort)
-}
-
 type HTTPServerTestSuite struct {
 	suite.Suite
-	server *httptest.Server
 }
 
 func (suite *HTTPServerTestSuite) SetupSuite() {
-	suite.server = httptest.NewServer(nil)
+	paramtable.Init()
-	registerDefaults()
+	ServeHTTP()
 }
 
 func (suite *HTTPServerTestSuite) TearDownSuite() {
-	defer suite.server.Close()
+	defer server.Close()
+	metricsServer = nil
+}
+
+func (suite *HTTPServerTestSuite) TestGetHTTPAddr() {
+	suite.Equal(getHTTPAddr(), ":"+DefaultListenPort)
+	testPort := "9092"
+	os.Setenv(ListenPortEnvKey, testPort)
+	suite.Equal(getHTTPAddr(), ":"+testPort)
 }
 
 func (suite *HTTPServerTestSuite) TestDefaultLogHandler() {
@@ -74,12 +71,12 @@ func (suite *HTTPServerTestSuite) TestDefaultLogHandler() {
 	payload, err := json.Marshal(map[string]any{"level": "error"})
 	suite.Require().NoError(err)
 
-	url := suite.server.URL + "/log/level"
+	url := "http://localhost:" + DefaultListenPort + "/log/level"
 	req, err := http.NewRequest(http.MethodPut, url, bytes.NewBuffer(payload))
 	req.Header.Set("Content-Type", "application/json")
 	suite.Require().NoError(err)
 
-	client := suite.server.Client()
+	client := http.Client{}
 	resp, err := client.Do(req)
 	suite.Require().NoError(err)
 	defer resp.Body.Close()
@@ -91,8 +88,8 @@ func (suite *HTTPServerTestSuite) TestDefaultLogHandler() {
 }
 
 func (suite *HTTPServerTestSuite) TestHealthzHandler() {
-	url := suite.server.URL + "/healthz"
+	url := "http://localhost:" + DefaultListenPort + "/healthz"
-	client := suite.server.Client()
+	client := http.Client{}
 
 	healthz.Register(&MockIndicator{"m1", commonpb.StateCode_Healthy})
 
@@ -121,6 +118,71 @@ func (suite *HTTPServerTestSuite) TestHealthzHandler() {
 	suite.Equal("{\"state\":\"component m2 state is Abnormal\",\"detail\":[{\"name\":\"m1\",\"code\":1},{\"name\":\"m2\",\"code\":2}]}", string(body))
 }
 
+func (suite *HTTPServerTestSuite) TestEventlogHandler() {
+	url := "http://localhost:" + DefaultListenPort + EventLogRouterPath
+	client := http.Client{}
+	req, _ := http.NewRequest(http.MethodGet, url, nil)
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := client.Do(req)
+	suite.Nil(err)
+	defer resp.Body.Close()
+	body, _ := io.ReadAll(resp.Body)
+	suite.True(strings.HasPrefix(string(body), "{\"status\":200,\"port\":"))
+}
+
+func (suite *HTTPServerTestSuite) TestPprofHandler() {
+	client := http.Client{}
+	testCases := []struct {
+		enable     bool
+		path       string
+		statusCode int
+		resp       []byte
+	}{
+		{true, "/debug/pprof/<script>scripty<script>", http.StatusNotFound, []byte("Unknown profile\n")},
+		{true, "/debug/pprof/heap", http.StatusOK, nil},
+		{true, "/debug/pprof/heap?debug=1", http.StatusOK, nil},
+		{true, "/debug/pprof/cmdline", http.StatusOK, nil},
+		{true, "/debug/pprof/profile?seconds=1", http.StatusOK, nil},
+		{true, "/debug/pprof/symbol", http.StatusOK, nil},
+		{true, "/debug/pprof/trace", http.StatusOK, nil},
+		{true, "/debug/pprof/mutex", http.StatusOK, nil},
+		{true, "/debug/pprof/block?seconds=1", http.StatusOK, nil},
+		{true, "/debug/pprof/goroutine?seconds=1", http.StatusOK, nil},
+		{true, "/debug/pprof/", http.StatusOK, []byte("Types of profiles available:")},
+		{false, "/debug/pprof/<script>scripty<script>", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/heap", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/heap?debug=1", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/cmdline", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/profile?seconds=1", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/symbol", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/trace", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/mutex", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/block?seconds=1", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/goroutine?seconds=1", http.StatusNotFound, []byte("404 page not found\n")},
+		{false, "/debug/pprof/", http.StatusNotFound, []byte("404 page not found\n")},
+	}
+	for _, tc := range testCases {
+		if tc.enable != paramtable.Get().HTTPCfg.EnablePprof.GetAsBool() {
+			continue
+		}
+		req, _ := http.NewRequest(http.MethodGet, "http://localhost:"+DefaultListenPort+tc.path, nil)
+		req.Header.Set("Content-Type", "application/json")
+		resp, err := client.Do(req)
+		suite.Nil(err)
+		if err == nil {
+			defer resp.Body.Close()
+			suite.Equal(tc.statusCode, resp.StatusCode)
+			body, err := io.ReadAll(resp.Body)
+			suite.Nil(err)
+			if resp.StatusCode != http.StatusOK {
+				suite.True(bytes.Equal(tc.resp, body))
+			}
+		} else {
+			fmt.Println(err.Error())
+		}
+	}
+}
+
 func TestHTTPServerSuite(t *testing.T) {
 	suite.Run(t, new(HTTPServerTestSuite))
 }

pkg/util/paramtable/http_param.go

@@ -5,6 +5,8 @@ type httpConfig struct {
 	DebugMode            ParamItem `refreshable:"false"`
 	Port                 ParamItem `refreshable:"false"`
 	AcceptTypeAllowInt64 ParamItem `refreshable:"false"`
+	EnablePprof          ParamItem `refreshable:"false"`
+	RequestTimeoutMs     ParamItem `refreshable:"false"`
 }
 
 func (p *httpConfig) init(base *BaseTable) {
@@ -44,4 +46,13 @@ func (p *httpConfig) init(base *BaseTable) {
 		Export:       true,
 	}
 	p.AcceptTypeAllowInt64.Init(base.mgr)
+
+	p.EnablePprof = ParamItem{
+		Key:          "proxy.http.enablePprof",
+		DefaultValue: "true",
+		Version:      "2.3.3",
+		Doc:          "Whether to enable pprof middleware on the metrics port",
+		Export:       true,
+	}
+	p.EnablePprof.Init(base.mgr)
 }

pkg/util/paramtable/http_param_test.go

@@ -14,4 +14,5 @@ func TestHTTPConfig_Init(t *testing.T) {
 	assert.Equal(t, cfg.DebugMode.GetAsBool(), false)
 	assert.Equal(t, cfg.Port.GetValue(), "")
 	assert.Equal(t, cfg.AcceptTypeAllowInt64.GetValue(), "false")
+	assert.Equal(t, cfg.EnablePprof.GetAsBool(), true)
 }