github.com/manicqin/nomad@v0.9.5/command/agent/pprof/pprof.go (about) 1 // Package profile is meant to be a near identical implemenation of 2 // https://golang.org/src/net/http/pprof/pprof.go 3 // It's purpose is to provide a way to accommodate the RPC endpoint style 4 // we use instead of traditional http handlers. 5 6 package pprof 7 8 import ( 9 "bytes" 10 "context" 11 "fmt" 12 "os" 13 "runtime" 14 "runtime/pprof" 15 "runtime/trace" 16 "strings" 17 "time" 18 ) 19 20 type ReqType string 21 22 const ( 23 CmdReq ReqType = "cmdline" 24 CPUReq ReqType = "cpu" 25 TraceReq ReqType = "trace" 26 LookupReq ReqType = "lookup" 27 28 ErrProfileNotFoundPrefix = "Pprof profile not found profile:" 29 ) 30 31 // NewErrProfileNotFound returns a new error caused by a pprof.Lookup 32 // profile not being found 33 func NewErrProfileNotFound(profile string) error { 34 return fmt.Errorf("%s %s", ErrProfileNotFoundPrefix, profile) 35 } 36 37 // IsErrProfileNotFound returns whether the error is due to a pprof profile 38 // being invalid 39 func IsErrProfileNotFound(err error) bool { 40 return err != nil && strings.Contains(err.Error(), ErrProfileNotFoundPrefix) 41 } 42 43 // Cmdline responds with the running program's 44 // command line, with arguments separated by NUL bytes. 45 func Cmdline() ([]byte, map[string]string, error) { 46 var buf bytes.Buffer 47 fmt.Fprintf(&buf, strings.Join(os.Args, "\x00")) 48 49 return buf.Bytes(), 50 map[string]string{ 51 "X-Content-Type-Options": "nosniff", 52 "Content-Type": "text/plain; charset=utf-8", 53 }, nil 54 } 55 56 // Profile generates a pprof.Profile report for the given profile name 57 // see runtime/pprof/pprof.go for available profiles. 58 func Profile(profile string, debug, gc int) ([]byte, map[string]string, error) { 59 p := pprof.Lookup(profile) 60 if p == nil { 61 return nil, nil, NewErrProfileNotFound(profile) 62 } 63 64 if profile == "heap" && gc > 0 { 65 runtime.GC() 66 } 67 68 var buf bytes.Buffer 69 if err := p.WriteTo(&buf, debug); err != nil { 70 return nil, nil, err 71 } 72 73 headers := map[string]string{ 74 "X-Content-Type-Options": "nosniff", 75 } 76 if debug != 0 { 77 headers["Content-Type"] = "text/plain; charset=utf-8" 78 } else { 79 headers["Content-Type"] = "application/octet-stream" 80 headers["Content-Disposition"] = fmt.Sprintf(`attachment; filename="%s"`, profile) 81 } 82 return buf.Bytes(), headers, nil 83 } 84 85 // CPUProfile generates a CPU Profile for a given duration 86 func CPUProfile(ctx context.Context, sec int) ([]byte, map[string]string, error) { 87 if sec <= 0 { 88 sec = 1 89 } 90 91 var buf bytes.Buffer 92 if err := pprof.StartCPUProfile(&buf); err != nil { 93 return nil, nil, err 94 } 95 96 sleep(ctx, time.Duration(sec)*time.Second) 97 98 pprof.StopCPUProfile() 99 100 return buf.Bytes(), 101 map[string]string{ 102 "X-Content-Type-Options": "nosniff", 103 "Content-Type": "application/octet-stream", 104 "Content-Disposition": `attachment; filename="profile"`, 105 }, nil 106 } 107 108 // Trace runs a trace profile for a given duration 109 func Trace(ctx context.Context, sec int) ([]byte, map[string]string, error) { 110 if sec <= 0 { 111 sec = 1 112 } 113 114 var buf bytes.Buffer 115 if err := trace.Start(&buf); err != nil { 116 return nil, nil, err 117 } 118 119 sleep(ctx, time.Duration(sec)*time.Second) 120 121 trace.Stop() 122 123 return buf.Bytes(), 124 map[string]string{ 125 "X-Content-Type-Options": "nosniff", 126 "Content-Type": "application/octet-stream", 127 "Content-Disposition": `attachment; filename="trace"`, 128 }, nil 129 } 130 131 func sleep(ctx context.Context, d time.Duration) { 132 // Sleep until duration is met or ctx is cancelled 133 select { 134 case <-time.After(d): 135 case <-ctx.Done(): 136 } 137 }