pprof.go

package pprof

import (
	"fmt"
	"net"
	"net/http"
	"net/http/pprof"
	"runtime"
	rpprof "runtime/pprof"
	"runtime/trace"
	"time"

	"github.com/gorilla/mux"
)

// Listen starts listening to the give address for fetch go profiling stats.
func Listen(addr string) (net.Listener, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return ln, err
	}
	runtime.SetBlockProfileRate(1)
	runtime.SetMutexProfileFraction(1)

	root := mux.NewRouter()
	sub := root.PathPrefix("/debug/pprof").Subrouter()
	sub.HandleFunc("/", pprof.Index)
	sub.HandleFunc("/cmdline", pprof.Cmdline)
	sub.HandleFunc("/profile", pprof.Profile)
	sub.HandleFunc("/symbol", pprof.Symbol)
	sub.HandleFunc("/trace", func(w http.ResponseWriter, r *http.Request) {
		d, err := time.ParseDuration(r.URL.Query().Get("d"))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(fmt.Sprintf("error parsing duration (?d=<duration>): %v", err)))
			return
		}

		if err := trace.Start(w); err != nil {
			w.WriteHeader(http.StatusServiceUnavailable)
			w.Write([]byte(fmt.Sprintf("error starting trace: %v", err)))
			return
		}

		timer := time.NewTimer(d)
		select {
		case <-r.Context().Done():
			timer.Stop()
		case <-timer.C:
		}

		trace.Stop()
	})

	for _, p := range rpprof.Profiles() {
		sub.Handle(fmt.Sprintf("/%s", p.Name()), pprof.Handler(p.Name()))
	}
	srv := &http.Server{
		Addr:    addr,
		Handler: root,
	}
	go srv.Serve(ln)
	return ln, nil
}