pprof/cmd/webpprof/main.go

package main

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"

	"github.com/gin-gonic/gin"
	"github.com/google/pprof/pkg/driver"
	"github.com/google/pprof/profile"
	"github.com/sirupsen/logrus"
)

func main() {
	r := gin.Default()

	r.Static("/html", "html")

	r.GET("/", func(c *gin.Context) { c.Redirect(http.StatusFound, "/html/") })
	r.POST("/upload", upload)
	r.GET("/pprof/:hash/:method", mapping)

	r.Run(":8090")
}

func mapping(c *gin.Context) {
	hash := c.Param("hash")
	if hash == "" {
		c.AbortWithStatus(http.StatusNotFound)
		return
	}

	filename := fmt.Sprintf("%s.pprof", hash)
	f, err := os.OpenFile(filename, os.O_RDONLY, os.ModePerm)
	if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}
	defer f.Close()
	prof, err := profile.Parse(f)
	if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	ui, err := driver.MakeWebInterface(prof, nil)
	if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	method := c.Param("method")
	switch method {
	case "", "dot":
		ui.Dot(c.Writer, c.Request)
	case "top":
		ui.Top(c.Writer, c.Request)
	case "disasm":
		ui.Disasm(c.Writer, c.Request)
	case "source":
		ui.Source(c.Writer, c.Request)
	case "peek":
		ui.Peek(c.Writer, c.Request)
	case "flamegraph":
		ui.Flamegraph(c.Writer, c.Request)
	case "saveconfig":
		ui.SaveConfig(c.Writer, c.Request)
	case "deleteconfig":
		ui.DeleteConfig(c.Writer, c.Request)
	default:
		c.AbortWithStatus(http.StatusNotFound)
	}
}

func upload(c *gin.Context) {
	f, _, err := c.Request.FormFile("file")
	if err != nil {
		logrus.Errorln("upload file error:", err)
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	fdata, _ := ioutil.ReadAll(f)
	f.Close()

	rhash := hash(fdata)

	if fileExist(rhash) { // 文件已存在
		c.Redirect(http.StatusFound, fmt.Sprintf("/pprof/%s/flamegraph", rhash))
		return
	}

	fileDst := fmt.Sprintf("%s.pprof", rhash)
	file, err := os.Create(fileDst)
	if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}
	defer file.Close()

	if _, err := file.Write(fdata); err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	c.JSONP(http.StatusOK, map[string]interface{}{
		"code": 0,
		"msg":  "ok",
		"data": rhash,
	})
}

func hash(data []byte) string {
	h := md5.New()
	h.Write(data)

	return hex.EncodeToString(h.Sum(nil))
}

func fileExist(file string) bool {
	return false
}