open-nomad/command/agent/pprof/pprof.go
Drew Bailey 45210ed901
Rename profile package to pprof
Address pr feedback, rename profile package to pprof to more accurately
describe its purpose. Adds gc param for heap lookup profiles.
2020-01-09 15:15:10 -05:00

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():
}
}