45210ed901
Address pr feedback, rename profile package to pprof to more accurately describe its purpose. Adds gc param for heap lookup profiles.
138 lines
3.3 KiB
Go
138 lines
3.3 KiB
Go
// Package profile is meant to be a near identical implemenation of
|
|
// https://golang.org/src/net/http/pprof/pprof.go
|
|
// It's purpose is to provide a way to accommodate the RPC endpoint style
|
|
// we use instead of traditional http handlers.
|
|
|
|
package pprof
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"runtime"
|
|
"runtime/pprof"
|
|
"runtime/trace"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type ReqType string
|
|
|
|
const (
|
|
CmdReq ReqType = "cmdline"
|
|
CPUReq ReqType = "cpu"
|
|
TraceReq ReqType = "trace"
|
|
LookupReq ReqType = "lookup"
|
|
|
|
ErrProfileNotFoundPrefix = "Pprof profile not found profile:"
|
|
)
|
|
|
|
// NewErrProfileNotFound returns a new error caused by a pprof.Lookup
|
|
// profile not being found
|
|
func NewErrProfileNotFound(profile string) error {
|
|
return fmt.Errorf("%s %s", ErrProfileNotFoundPrefix, profile)
|
|
}
|
|
|
|
// IsErrProfileNotFound returns whether the error is due to a pprof profile
|
|
// being invalid
|
|
func IsErrProfileNotFound(err error) bool {
|
|
return err != nil && strings.Contains(err.Error(), ErrProfileNotFoundPrefix)
|
|
}
|
|
|
|
// Cmdline responds with the running program's
|
|
// command line, with arguments separated by NUL bytes.
|
|
func Cmdline() ([]byte, map[string]string, error) {
|
|
var buf bytes.Buffer
|
|
fmt.Fprintf(&buf, strings.Join(os.Args, "\x00"))
|
|
|
|
return buf.Bytes(),
|
|
map[string]string{
|
|
"X-Content-Type-Options": "nosniff",
|
|
"Content-Type": "text/plain; charset=utf-8",
|
|
}, nil
|
|
}
|
|
|
|
// Profile generates a pprof.Profile report for the given profile name
|
|
// see runtime/pprof/pprof.go for available profiles.
|
|
func Profile(profile string, debug, gc int) ([]byte, map[string]string, error) {
|
|
p := pprof.Lookup(profile)
|
|
if p == nil {
|
|
return nil, nil, NewErrProfileNotFound(profile)
|
|
}
|
|
|
|
if profile == "heap" && gc > 0 {
|
|
runtime.GC()
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := p.WriteTo(&buf, debug); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
headers := map[string]string{
|
|
"X-Content-Type-Options": "nosniff",
|
|
}
|
|
if debug != 0 {
|
|
headers["Content-Type"] = "text/plain; charset=utf-8"
|
|
} else {
|
|
headers["Content-Type"] = "application/octet-stream"
|
|
headers["Content-Disposition"] = fmt.Sprintf(`attachment; filename="%s"`, profile)
|
|
}
|
|
return buf.Bytes(), headers, nil
|
|
}
|
|
|
|
// CPUProfile generates a CPU Profile for a given duration
|
|
func CPUProfile(ctx context.Context, sec int) ([]byte, map[string]string, error) {
|
|
if sec <= 0 {
|
|
sec = 1
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := pprof.StartCPUProfile(&buf); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
sleep(ctx, time.Duration(sec)*time.Second)
|
|
|
|
pprof.StopCPUProfile()
|
|
|
|
return buf.Bytes(),
|
|
map[string]string{
|
|
"X-Content-Type-Options": "nosniff",
|
|
"Content-Type": "application/octet-stream",
|
|
"Content-Disposition": `attachment; filename="profile"`,
|
|
}, nil
|
|
}
|
|
|
|
// Trace runs a trace profile for a given duration
|
|
func Trace(ctx context.Context, sec int) ([]byte, map[string]string, error) {
|
|
if sec <= 0 {
|
|
sec = 1
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := trace.Start(&buf); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
sleep(ctx, time.Duration(sec)*time.Second)
|
|
|
|
trace.Stop()
|
|
|
|
return buf.Bytes(),
|
|
map[string]string{
|
|
"X-Content-Type-Options": "nosniff",
|
|
"Content-Type": "application/octet-stream",
|
|
"Content-Disposition": `attachment; filename="trace"`,
|
|
}, nil
|
|
}
|
|
|
|
func sleep(ctx context.Context, d time.Duration) {
|
|
// Sleep until duration is met or ctx is cancelled
|
|
select {
|
|
case <-time.After(d):
|
|
case <-ctx.Done():
|
|
}
|
|
}
|