github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/ginx/ginpprof/pprof.go (about)

     1  package ginpprof
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"net/http/pprof"
     7  	"strings"
     8  
     9  	"github.com/gin-gonic/gin"
    10  )
    11  
    12  // Wrap adds several routes from package `net/http/pprof` to *gin.Engine object
    13  func Wrap(router interface{}) {
    14  	if r, ok := router.(*gin.Engine); ok {
    15  		WrapGroup(&r.RouterGroup)
    16  	} else if r, ok := router.(*gin.RouterGroup); ok {
    17  		WrapGroup(r)
    18  	} else if r, ok := router.(*http.ServeMux); ok {
    19  		WrapServeMux(r)
    20  	} else {
    21  		panic(fmt.Errorf("please wrap *gin.Engine or *gin.RouterGroup"))
    22  	}
    23  }
    24  
    25  // WrapGroup adds several routes from package `net/http/pprof` to *gin.RouterGroup object
    26  func WrapGroup(router *gin.RouterGroup) {
    27  	basePath := strings.TrimSuffix(router.BasePath(), "/")
    28  
    29  	var prefix string
    30  
    31  	switch {
    32  	case basePath == "":
    33  		prefix = ""
    34  	case strings.HasSuffix(basePath, "/debug"):
    35  		prefix = "/debug"
    36  	case strings.HasSuffix(basePath, "/debug/pprof"):
    37  		prefix = "/debug/pprof"
    38  	}
    39  
    40  	for _, r := range routers {
    41  		router.Handle(r.Method, strings.TrimPrefix(r.Path, prefix), r.Handler)
    42  	}
    43  }
    44  
    45  // WrapServeMux adds several routes from package `net/http/pprof` to *http.ServeMux object
    46  func WrapServeMux(mux *http.ServeMux) {
    47  	mux.HandleFunc("/debug/pprof/", pprof.Index)
    48  	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
    49  	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
    50  	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
    51  	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
    52  
    53  	mux.Handle("/debug/pprof/heap", pprof.Handler("heap"))
    54  	mux.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
    55  	mux.Handle("/debug/pprof/allocs", pprof.Handler("allocs"))
    56  	mux.Handle("/debug/pprof/block", pprof.Handler("block"))
    57  	mux.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
    58  	mux.Handle("/debug/pprof/mutex", pprof.Handler("mutex"))
    59  }
    60  
    61  // PprofServeHTTP servces HTTP requests for /debug/pprof requests.
    62  func PprofServeHTTP(w http.ResponseWriter, r *http.Request) bool {
    63  	switch r.URL.Path {
    64  	case "/debug/pprof":
    65  		pprof.Index(w, r)
    66  	case "/debug/pprof/cmdline":
    67  		pprof.Cmdline(w, r)
    68  	case "/debug/pprof/symbol":
    69  		pprof.Symbol(w, r)
    70  	case "/debug/pprof/profile":
    71  		pprof.Profile(w, r)
    72  	case "/debug/pprof/trace":
    73  		pprof.Trace(w, r)
    74  	case "/debug/pprof/heap":
    75  		pprof.Handler("heap").ServeHTTP(w, r)
    76  	case "/debug/pprof/goroutine":
    77  		pprof.Handler("goroutine").ServeHTTP(w, r)
    78  	case "/debug/pprof/allocs":
    79  		pprof.Handler("allocs").ServeHTTP(w, r)
    80  	case "/debug/pprof/block":
    81  		pprof.Handler("block").ServeHTTP(w, r)
    82  	case "/debug/pprof/threadcreate":
    83  		pprof.Handler("threadcreate").ServeHTTP(w, r)
    84  	case "/debug/pprof/mutex":
    85  		pprof.Handler("mutex").ServeHTTP(w, r)
    86  	default:
    87  		return false
    88  	}
    89  
    90  	return true
    91  }
    92  
    93  var routers = []struct {
    94  	Method  string
    95  	Path    string
    96  	Handler gin.HandlerFunc
    97  }{
    98  	{"GET", "/debug/pprof/", func(c *gin.Context) {
    99  		pprof.Index(c.Writer, c.Request)
   100  	}},
   101  	{"GET", "/debug/pprof/cmdline", func(c *gin.Context) {
   102  		pprof.Cmdline(c.Writer, c.Request)
   103  	}},
   104  	{"GET", "/debug/pprof/profile", func(c *gin.Context) {
   105  		pprof.Profile(c.Writer, c.Request)
   106  	}},
   107  	{"GET", "/debug/pprof/symbol", func(c *gin.Context) {
   108  		pprof.Symbol(c.Writer, c.Request)
   109  	}},
   110  	{"POST", "/debug/pprof/symbol", func(c *gin.Context) {
   111  		pprof.Symbol(c.Writer, c.Request)
   112  	}},
   113  	{"GET", "/debug/pprof/trace", func(c *gin.Context) {
   114  		pprof.Trace(c.Writer, c.Request)
   115  	}},
   116  	{"GET", "/debug/pprof/heap", func(c *gin.Context) {
   117  		pprof.Handler("heap").ServeHTTP(c.Writer, c.Request)
   118  	}},
   119  	{"GET", "/debug/pprof/goroutine", func(c *gin.Context) {
   120  		pprof.Handler("goroutine").ServeHTTP(c.Writer, c.Request)
   121  	}},
   122  	{"GET", "/debug/pprof/allocs", func(c *gin.Context) {
   123  		pprof.Handler("allocs").ServeHTTP(c.Writer, c.Request)
   124  	}},
   125  	{"GET", "/debug/pprof/block", func(c *gin.Context) {
   126  		pprof.Handler("block").ServeHTTP(c.Writer, c.Request)
   127  	}},
   128  	{"GET", "/debug/pprof/threadcreate", func(c *gin.Context) {
   129  		pprof.Handler("threadcreate").ServeHTTP(c.Writer, c.Request)
   130  	}},
   131  	{"GET", "/debug/pprof/mutex", func(c *gin.Context) {
   132  		pprof.Handler("mutex").ServeHTTP(c.Writer, c.Request)
   133  	}},
   134  }