github.com/hernad/nomad@v1.6.112/command/agent/pprof/pprof.go (about)

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