diff --git a/internal/http/server.go b/internal/http/server.go index cbd483f3ac..cb5e50521b 100644 --- a/internal/http/server.go +++ b/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)) } diff --git a/internal/http/server_test.go b/internal/http/server_test.go index 7cba7885d0..852353dfd3 100644 --- a/internal/http/server_test.go +++ b/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) - registerDefaults() + paramtable.Init() + 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" - client := suite.server.Client() + url := "http://localhost:" + DefaultListenPort + "/healthz" + 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/