diff --git a/go.mod b/go.mod index 1336d502ee2..ee121d16865 100644 --- a/go.mod +++ b/go.mod @@ -37,6 +37,7 @@ require ( github.com/pingcap/tidb-dashboard v0.0.0-20240327052925-0f035e0e22ee github.com/prometheus/client_golang v1.11.1 github.com/prometheus/common v0.26.0 + github.com/rhysh/autoprof v0.0.0-20240708032903-aed43a1dedec github.com/sasha-s/go-deadlock v0.2.0 github.com/spf13/cobra v1.0.0 github.com/spf13/pflag v1.0.5 diff --git a/go.sum b/go.sum index 38979e85bb2..6359965fe19 100644 --- a/go.sum +++ b/go.sum @@ -428,6 +428,8 @@ github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O github.com/prometheus/procfs v0.6.0 h1:mxy4L2jP6qMonqmq+aTtOx1ifVWUgG/TAmntgbh3xv4= github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA= github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU= +github.com/rhysh/autoprof v0.0.0-20240708032903-aed43a1dedec h1:EeYHkttFMOX34U4NBrZOKk1+I0vW5chptpdTvC1i9Hc= +github.com/rhysh/autoprof v0.0.0-20240708032903-aed43a1dedec/go.mod h1:4FBop5l5Ur0n+mQ7rMX7XMdfhewY7N2aB/OjAlwFfuE= github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg= github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ= github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4= diff --git a/server/api/router.go b/server/api/router.go index cf604c47738..96a7ffb1f8b 100644 --- a/server/api/router.go +++ b/server/api/router.go @@ -15,14 +15,19 @@ package api import ( + "bytes" + "context" + "fmt" "net/http" "net/http/pprof" "reflect" "runtime" "strings" + "time" "github.com/gorilla/mux" "github.com/pingcap/failpoint" + "github.com/rhysh/autoprof" "github.com/tikv/pd/pkg/apiutil" "github.com/tikv/pd/pkg/audit" "github.com/tikv/pd/pkg/ratelimit" @@ -343,6 +348,7 @@ func createRouter(prefix string, svr *server.Server) *mux.Router { registerFunc(apiRouter, "/debug/pprof/goroutine", pprofHandler.PProfGoroutine) registerFunc(apiRouter, "/debug/pprof/threadcreate", pprofHandler.PProfThreadcreate) registerFunc(apiRouter, "/debug/pprof/zip", pprofHandler.PProfZip) + registerFunc(apiRouter, "/debug/pprof/autoprof", collectHandler) // service GC safepoint API serviceGCSafepointHandler := newServiceGCSafepointHandler(svr, rd) @@ -395,3 +401,32 @@ func createRouter(prefix string, svr *server.Server) *mux.Router { return rootRouter } + +func collectHandler(w http.ResponseWriter, r *http.Request) { + meta := autoprof.CurrentArchiveMeta() + opt := &autoprof.ArchiveOptions{ + CPUProfileDuration: 3 * time.Second, + ExecutionTraceDuration: 3 * time.Second, + } + + buf, err := collect(r.Context(), meta, opt) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/zip") + w.Header().Set("Content-Disposition", "attachment; filename=autoprof.zip") + if _, err := w.Write(buf.Bytes()); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } +} + +func collect(ctx context.Context, meta *autoprof.ArchiveMeta, opt *autoprof.ArchiveOptions) (bytes.Buffer, error) { + var buf bytes.Buffer + + err := autoprof.NewZipCollector(&buf, meta, opt).Run(context.Background()) + if err != nil { + return buf, fmt.Errorf("autoprof.NewZipCollector.Run: %w", err) + } + return buf, nil +}