server: move all routes to mux router and make sure pprof paths are g…

…ood (pingcap#52236) close pingcap#51914
Leavrth · Mar 29, 2024 · 208df64 · 208df64
1 parent 8bba826
commit 208df64
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 13 deletions.
diff --git a/pkg/server/handler/tests/BUILD.bazel b/pkg/server/handler/tests/BUILD.bazel
@@ -9,7 +9,7 @@ go_test(
         "main_test.go",
     ],
     flaky = True,
-    shard_count = 37,
+    shard_count = 38,
     deps = [
         "//pkg/config",
         "//pkg/ddl",

diff --git a/pkg/server/handler/tests/http_handler_serial_test.go b/pkg/server/handler/tests/http_handler_serial_test.go
@@ -440,6 +440,37 @@ func TestTiFlashReplica(t *testing.T) {
 	checkFunc()
 }
 
+func TestDebugRoutes(t *testing.T) {
+	ts := createBasicHTTPHandlerTestSuite()
+	ts.startServer(t)
+	defer ts.stopServer(t)
+
+	debugRoutes := []string{
+		"/debug/pprof/",
+		"/debug/pprof/heap?debug=1",
+		"/debug/pprof/goroutine?debug=1",
+		"/debug/pprof/goroutine?debug=2",
+		"/debug/pprof/allocs?debug=1",
+		"/debug/pprof/block?debug=1",
+		"/debug/pprof/threadcreate?debug=1",
+		"/debug/pprof/cmdline",
+		"/debug/pprof/profile",
+		"/debug/pprof/mutex?debug=1",
+		"/debug/pprof/symbol",
+		"/debug/pprof/trace",
+		"/debug/pprof/profile",
+		"/debug/gogc",
+		// "/debug/zip", // this creates unexpected goroutines which will make goleak complain, so we skip it for now
+		"/debug/ballast-object-sz",
+	}
+	for _, route := range debugRoutes {
+		resp, err := ts.FetchStatus(route)
+		require.NoError(t, err, fmt.Sprintf("GET route %s failed", route))
+		require.Equal(t, http.StatusOK, resp.StatusCode, fmt.Sprintf("GET route %s failed", route))
+		require.NoError(t, resp.Body.Close())
+	}
+}
+
 func TestFailpointHandler(t *testing.T) {
 	ts := createBasicHTTPHandlerTestSuite()
 

diff --git a/pkg/server/http_status.go b/pkg/server/http_status.go
@@ -294,14 +294,12 @@ func (s *Server) startHTTPServer() {
 		router.PathPrefix("/static/").Handler(http.StripPrefix("/static", http.FileServer(static.Data)))
 	}
 
-	serverMux := http.NewServeMux()
-	serverMux.Handle("/", router)
-
-	serverMux.HandleFunc("/debug/pprof/", pprof.Index)
-	serverMux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
-	serverMux.HandleFunc("/debug/pprof/profile", cpuprofile.ProfileHTTPHandler)
-	serverMux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
-	serverMux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	router.HandleFunc("/debug/pprof/profile", cpuprofile.ProfileHTTPHandler)
+	router.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	router.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	// Other /debug/pprof paths not covered above are redirected to pprof.Index.
+	router.PathPrefix("/debug/pprof/").HandlerFunc(pprof.Index)
 
 	ballast := newBallast(s.cfg.MaxBallastObjectSize)
 	{
@@ -310,9 +308,9 @@ func (s *Server) startHTTPServer() {
 			logutil.BgLogger().Error("set initial ballast object size failed", zap.Error(err))
 		}
 	}
-	serverMux.HandleFunc("/debug/ballast-object-sz", ballast.GenHTTPHandler())
+	router.HandleFunc("/debug/ballast-object-sz", ballast.GenHTTPHandler())
 
-	serverMux.HandleFunc("/debug/gogc", func(w http.ResponseWriter, r *http.Request) {
+	router.HandleFunc("/debug/gogc", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodGet:
 			_, err := w.Write([]byte(strconv.Itoa(util.GetGOGC())))
@@ -337,7 +335,7 @@ func (s *Server) startHTTPServer() {
 		}
 	})
 
-	serverMux.HandleFunc("/debug/zip", func(w http.ResponseWriter, r *http.Request) {
+	router.HandleFunc("/debug/zip", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="tidb_debug"`+time.Now().Format("20060102150405")+".zip"))
 
 		// dump goroutine/heap/mutex
@@ -423,7 +421,7 @@ func (s *Server) startHTTPServer() {
 
 	// failpoint is enabled only for tests so we can add some http APIs here for tests.
 	failpoint.Inject("enableTestAPI", func() {
-		serverMux.HandleFunc("/fail/", func(w http.ResponseWriter, r *http.Request) {
+		router.PathPrefix("/fail/").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			r.URL.Path = strings.TrimPrefix(r.URL.Path, "/fail")
 			new(failpoint.HttpHandler).ServeHTTP(w, r)
 		})
@@ -467,6 +465,9 @@ func (s *Server) startHTTPServer() {
 			logutil.BgLogger().Error("write HTTP index page failed", zap.Error(err))
 		}
 	})
+
+	serverMux := http.NewServeMux()
+	serverMux.Handle("/", router)
 	s.startStatusServerAndRPCServer(serverMux)
 }