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  }