test pprof headers and profile methods

tidy up, add comments

clean up seconds param assignment
This commit is contained in:
Drew Bailey 2019-12-11 14:58:41 -05:00
parent fc37448683
commit 3672414888
No known key found for this signature in database
GPG Key ID: FBA61B9FB7CCE1A7
7 changed files with 243 additions and 32 deletions

View File

@ -39,18 +39,19 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP
var resp []byte
var err error
var headers map[string]string
// Determine which profile to run
// and generate profile. Blocks for args.Seconds
switch args.ReqType {
case profile.CPUReq:
resp, err = profile.CPUProfile(context.TODO(), args.Seconds)
resp, headers, err = profile.CPUProfile(context.TODO(), args.Seconds)
case profile.CmdReq:
resp, err = profile.Cmdline()
resp, headers, err = profile.Cmdline()
case profile.LookupReq:
resp, err = profile.Profile(args.Profile, args.Debug)
resp, headers, err = profile.Profile(args.Profile, args.Debug)
case profile.TraceReq:
resp, err = profile.Trace(context.TODO(), args.Seconds)
resp, headers, err = profile.Trace(context.TODO(), args.Seconds)
}
if err != nil {
@ -63,6 +64,7 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP
// Copy profile response to reply
reply.Payload = resp
reply.AgentID = a.c.NodeID()
reply.HTTPHeaders = headers
return nil
}

View File

@ -85,6 +85,8 @@ type AgentPprofResponse struct {
// Payload is the generated pprof profile
Payload []byte
HTTPHeaders map[string]string
}
// AllocFileInfo holds information about a file inside the AllocDir

View File

@ -334,7 +334,7 @@ func (s *HTTPServer) AgentForceLeaveRequest(resp http.ResponseWriter, req *http.
return nil, err
}
func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Request) ([]byte, error) {
path := strings.TrimPrefix(req.URL.Path, "/v1/agent/pprof/")
switch {
case path == "":
@ -357,7 +357,7 @@ func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Reque
}
}
func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWriter, req *http.Request) (interface{}, error) {
func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWriter, req *http.Request) ([]byte, error) {
var secret string
s.parseToken(req, &secret)
@ -386,11 +386,10 @@ func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWrite
if secondsParam == "" {
seconds = 1
} else {
sec, err := strconv.Atoi(secondsParam)
seconds, err = strconv.Atoi(secondsParam)
if err != nil {
return nil, CodedError(400, err.Error())
}
seconds = sec
}
// Create the request
@ -435,11 +434,11 @@ func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWrite
return nil, CodedError(code, msg)
}
// Pprof cmdline is not a typical pprof
// so just return string instead of bytes
if args.ReqType == profile.CmdReq {
return string(reply.Payload), nil
// Set headers from profile request
for k, v := range reply.HTTPHeaders {
resp.Header().Set(k, v)
}
return reply.Payload, nil
}

View File

@ -186,7 +186,7 @@ func (s *HTTPServer) registerHandlers(enableDebug bool) {
s.mux.HandleFunc("/v1/agent/health", s.wrap(s.HealthRequest))
s.mux.HandleFunc("/v1/agent/monitor", s.wrap(s.AgentMonitor))
s.mux.HandleFunc("/v1/agent/pprof/", s.wrap(s.AgentPprofRequest))
s.mux.HandleFunc("/v1/agent/pprof/", s.wrapNonJSON(s.AgentPprofRequest))
s.mux.HandleFunc("/v1/metrics", s.wrap(s.MetricsRequest))
@ -358,6 +358,54 @@ func (s *HTTPServer) wrap(handler func(resp http.ResponseWriter, req *http.Reque
return f
}
// wrapNonJON is used to wrap functions returning non JSON
// serializeable data to make them more convenient. It is primarily
// responsible for setting nomad headers and logging.
func (s *HTTPServer) wrapNonJSON(handler func(resp http.ResponseWriter, req *http.Request) ([]byte, error)) func(resp http.ResponseWriter, req *http.Request) {
f := func(resp http.ResponseWriter, req *http.Request) {
setHeaders(resp, s.agent.config.HTTPAPIResponseHeaders)
// Invoke the handler
reqURL := req.URL.String()
start := time.Now()
defer func() {
s.logger.Debug("request complete", "method", req.Method, "path", reqURL, "duration", time.Now().Sub(start))
}()
obj, err := handler(resp, req)
// Check for an error
if err != nil {
code := 500
errMsg := err.Error()
if http, ok := err.(HTTPCodedError); ok {
code = http.Code()
} else if ecode, emsg, ok := structs.CodeFromRPCCodedErr(err); ok {
code = ecode
errMsg = emsg
} else {
// RPC errors get wrapped, so manually unwrap by only looking at their suffix
if strings.HasSuffix(errMsg, structs.ErrPermissionDenied.Error()) {
errMsg = structs.ErrPermissionDenied.Error()
code = 403
} else if strings.HasSuffix(errMsg, structs.ErrTokenNotFound.Error()) {
errMsg = structs.ErrTokenNotFound.Error()
code = 403
}
}
resp.WriteHeader(code)
resp.Write([]byte(errMsg))
s.logger.Error("request failed", "method", req.Method, "path", reqURL, "error", err, "code", code)
return
}
// write response
if obj != nil {
resp.Write(obj)
}
}
return f
}
// decodeBody is used to decode a JSON request body
func decodeBody(req *http.Request, out interface{}) error {
dec := json.NewDecoder(req.Body)

View File

@ -42,64 +42,86 @@ func IsErrProfileNotFound(err error) bool {
// Cmdline responds with the running program's
// command line, with arguments separated by NUL bytes.
func Cmdline() ([]byte, error) {
func Cmdline() ([]byte, map[string]string, error) {
var buf bytes.Buffer
fmt.Fprintf(&buf, strings.Join(os.Args, "\x00"))
return buf.Bytes(), nil
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 int) ([]byte, error) {
func Profile(profile string, debug int) ([]byte, map[string]string, error) {
p := pprof.Lookup(profile)
if p == nil {
return nil, NewErrProfileNotFound(profile)
return nil, nil, NewErrProfileNotFound(profile)
}
var buf bytes.Buffer
if err := p.WriteTo(&buf, debug); err != nil {
return nil, err
return nil, nil, err
}
return buf.Bytes(), nil
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, error) {
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 {
// trace.Start failed, no writes yet
return nil, err
return nil, nil, err
}
sleep(ctx, time.Duration(sec)*time.Second)
pprof.StopCPUProfile()
return buf.Bytes(), nil
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, error) {
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 {
// trace.Start failed, no writes yet
return nil, err
return nil, nil, err
}
sleep(context.TODO(), time.Duration(sec)*time.Second)
sleep(ctx, time.Duration(sec)*time.Second)
trace.Stop()
return buf.Bytes(), nil
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) {

View File

@ -0,0 +1,134 @@
package profile
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
func TestProfile(t *testing.T) {
cases := []struct {
desc string
profile string
debug int
expectedHeaders map[string]string
expectedErr error
}{
{
desc: "profile that exists",
profile: "goroutine",
expectedHeaders: map[string]string{
"X-Content-Type-Options": "nosniff",
"Content-Type": "application/octet-stream",
"Content-Disposition": `attachment; filename="goroutine"`,
},
},
{
desc: "profile that does not exist",
profile: "nonexistent",
expectedErr: NewErrProfileNotFound("nonexistent"),
expectedHeaders: nil,
},
{
desc: "profile with debug enabled",
profile: "allocs",
debug: 1,
expectedHeaders: map[string]string{
"X-Content-Type-Options": "nosniff",
"Content-Type": "text/plain; charset=utf-8",
},
},
}
for _, tc := range cases {
t.Run(tc.desc, func(t *testing.T) {
resp, headers, err := Profile(tc.profile, tc.debug)
require.Equal(t, tc.expectedHeaders, headers)
if tc.expectedErr != nil {
require.Nil(t, resp)
require.Equal(t, err, tc.expectedErr)
} else {
require.NotNil(t, resp)
}
})
}
}
func TestCPUProfile(t *testing.T) {
cases := []struct {
desc string
expectedHeaders map[string]string
}{
{
desc: "successful cpu profile",
expectedHeaders: map[string]string{
"X-Content-Type-Options": "nosniff",
"Content-Type": "application/octet-stream",
"Content-Disposition": `attachment; filename="profile"`,
},
},
}
for _, tc := range cases {
t.Run(tc.desc, func(t *testing.T) {
resp, headers, err := CPUProfile(context.Background(), 0)
require.NoError(t, err)
require.Equal(t, tc.expectedHeaders, headers)
require.NotNil(t, resp)
})
}
}
func TestTrace(t *testing.T) {
cases := []struct {
desc string
expectedHeaders map[string]string
}{
{
desc: "successful trace profile",
expectedHeaders: map[string]string{
"X-Content-Type-Options": "nosniff",
"Content-Type": "application/octet-stream",
"Content-Disposition": `attachment; filename="trace"`,
},
},
}
for _, tc := range cases {
t.Run(tc.desc, func(t *testing.T) {
resp, headers, err := Trace(context.Background(), 0)
require.NoError(t, err)
require.Equal(t, tc.expectedHeaders, headers)
require.NotNil(t, resp)
})
}
}
func TestCmdline(t *testing.T) {
cases := []struct {
desc string
expectedHeaders map[string]string
}{
{
desc: "successful cmdline request",
expectedHeaders: map[string]string{
"X-Content-Type-Options": "nosniff",
"Content-Type": "text/plain; charset=utf-8",
},
},
}
for _, tc := range cases {
t.Run(tc.desc, func(t *testing.T) {
resp, headers, err := Cmdline()
require.NoError(t, err)
require.Equal(t, tc.expectedHeaders, headers)
require.NotNil(t, resp)
})
}
}

View File

@ -62,21 +62,24 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP
// Process the request on this server
var resp []byte
var err error
var headers map[string]string
// Mark which server fulfilled the request
reply.AgentID = a.srv.serf.LocalMember().Name
// Determine which profile to run
// and generate profile. Blocks for args.Seconds
// Our RPC endpoints currently don't support context
// or request cancellation so stubbing with TODO
switch args.ReqType {
case profile.CPUReq:
resp, err = profile.CPUProfile(context.TODO(), args.Seconds)
resp, headers, err = profile.CPUProfile(context.TODO(), args.Seconds)
case profile.CmdReq:
resp, err = profile.Cmdline()
resp, headers, err = profile.Cmdline()
case profile.LookupReq:
resp, err = profile.Profile(args.Profile, args.Debug)
resp, headers, err = profile.Profile(args.Profile, args.Debug)
case profile.TraceReq:
resp, err = profile.Trace(context.TODO(), args.Seconds)
resp, headers, err = profile.Trace(context.TODO(), args.Seconds)
default:
err = structs.NewErrRPCCoded(404, "Unknown profile request type")
}
@ -90,6 +93,7 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP
// Copy profile response to reply
reply.Payload = resp
reply.HTTPHeaders = headers
return nil
}