pprofhack.go

package main

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/google/pprof/driver"
)

func main() {
	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	go func() {
		select {
		case <-ctx.Done():
			cancel()
			os.Exit(0)
		}
	}()

	profilePath := "/tmp/foo/heap.out"
	if err := spawnWebServer(ctx, profilePath, 8080); err != nil {
		panic(err)
	}

}

func spawnWebServer(ctx context.Context, file string, port int) error {
	return driver.PProf(&driver.Options{
		Flagset: &MockFlags{file, port},
		HTTPServer: func(args *driver.HTTPServerArgs) error {
			ln, err := new(net.ListenConfig).Listen(ctx, "tcp", args.Hostport)
			if err != nil {
				return err
			}
			handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
				h := args.Handlers[req.URL.Path]
				if h == nil {
					h = http.DefaultServeMux
				}
				h.ServeHTTP(w, req)
			})

			mux := http.NewServeMux()
			mux.Handle("/", handler)
			mux.Handle("/ui/", http.StripPrefix("/ui", handler))
			s := &http.Server{Handler: mux}
			return s.Serve(ln)
		},
	})
}

var _ driver.FlagSet = &MockFlags{}

type MockFlags struct {
	profilePath string
	httpPort    int
}

func (m *MockFlags) String(o, d, c string) *string {
	if o == "http" {
		d = fmt.Sprintf("localhost:%d", m.httpPort)
	}
	return &d
}

func (m *MockFlags) Parse(usage func()) []string {
	return []string{m.profilePath}
}

func (*MockFlags) Bool(o string, d bool, c string) *bool          { return &d }
func (*MockFlags) Int(o string, d int, c string) *int             { return &d }
func (*MockFlags) Float64(o string, d float64, c string) *float64 { return &d }
func (*MockFlags) StringList(o, d, c string) *[]*string           { return &[]*string{&d} }
func (f *MockFlags) ExtraUsage() string                           { return "" }
func (f *MockFlags) AddExtraUsage(eu string)                      {}

Context: google/pprof#753