github.com/palisadeinc/bor@v0.0.0-20230615125219-ab7196213d15/internal/cli/server/pprof/pprof.go (about)

     1  package pprof
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"net/http"
     8  	"runtime"
     9  	"runtime/pprof"
    10  	"runtime/trace"
    11  	"time"
    12  
    13  	"github.com/ethereum/go-ethereum/log"
    14  )
    15  
    16  // Profile generates a pprof.Profile report for the given profile name.
    17  func Profile(profile string, debug, gc int) ([]byte, map[string]string, error) {
    18  	p := pprof.Lookup(profile)
    19  	if p == nil {
    20  		return nil, nil, fmt.Errorf("profile '%s' not found", profile)
    21  	}
    22  
    23  	if profile == "heap" && gc > 0 {
    24  		runtime.GC()
    25  	}
    26  
    27  	var buf bytes.Buffer
    28  	if err := p.WriteTo(&buf, debug); err != nil {
    29  		return nil, nil, err
    30  	}
    31  
    32  	headers := map[string]string{
    33  		"X-Content-Type-Options": "nosniff",
    34  	}
    35  	if debug != 0 {
    36  		headers["Content-Type"] = "text/plain; charset=utf-8"
    37  	} else {
    38  		headers["Content-Type"] = "application/octet-stream"
    39  		headers["Content-Disposition"] = fmt.Sprintf(`attachment; filename="%s"`, profile)
    40  	}
    41  	return buf.Bytes(), headers, nil
    42  }
    43  
    44  // CPUProfile generates a CPU Profile for a given duration
    45  func CPUProfile(ctx context.Context, sec int) ([]byte, map[string]string, error) {
    46  	if sec <= 0 {
    47  		sec = 1
    48  	}
    49  
    50  	var buf bytes.Buffer
    51  	if err := pprof.StartCPUProfile(&buf); err != nil {
    52  		return nil, nil, err
    53  	}
    54  
    55  	sleep(ctx, time.Duration(sec)*time.Second)
    56  
    57  	pprof.StopCPUProfile()
    58  
    59  	return buf.Bytes(),
    60  		map[string]string{
    61  			"X-Content-Type-Options": "nosniff",
    62  			"Content-Type":           "application/octet-stream",
    63  			"Content-Disposition":    `attachment; filename="profile"`,
    64  		}, nil
    65  }
    66  
    67  // CPUProfile generates a CPU Profile for a given duration
    68  func CPUProfileWithChannel(done chan bool) ([]byte, map[string]string, error) {
    69  	var buf bytes.Buffer
    70  	if err := pprof.StartCPUProfile(&buf); err != nil {
    71  		return nil, nil, err
    72  	}
    73  
    74  	select {
    75  	case <-done:
    76  	case <-time.After(30 * time.Second):
    77  	}
    78  
    79  	pprof.StopCPUProfile()
    80  
    81  	return buf.Bytes(),
    82  		map[string]string{
    83  			"X-Content-Type-Options": "nosniff",
    84  			"Content-Type":           "application/octet-stream",
    85  			"Content-Disposition":    `attachment; filename="profile"`,
    86  		}, nil
    87  }
    88  
    89  // Trace runs a trace profile for a given duration
    90  func Trace(ctx context.Context, sec int) ([]byte, map[string]string, error) {
    91  	if sec <= 0 {
    92  		sec = 1
    93  	}
    94  
    95  	var buf bytes.Buffer
    96  	if err := trace.Start(&buf); err != nil {
    97  		return nil, nil, err
    98  	}
    99  
   100  	sleep(ctx, time.Duration(sec)*time.Second)
   101  
   102  	trace.Stop()
   103  
   104  	return buf.Bytes(),
   105  		map[string]string{
   106  			"X-Content-Type-Options": "nosniff",
   107  			"Content-Type":           "application/octet-stream",
   108  			"Content-Disposition":    `attachment; filename="trace"`,
   109  		}, nil
   110  }
   111  
   112  func sleep(ctx context.Context, d time.Duration) {
   113  	// Sleep until duration is met or ctx is cancelled
   114  	select {
   115  	case <-time.After(d):
   116  	case <-ctx.Done():
   117  	}
   118  }
   119  
   120  func SetMemProfileRate(rate int) {
   121  	runtime.MemProfileRate = rate
   122  }
   123  
   124  func SetSetBlockProfileRate(rate int) {
   125  	runtime.SetBlockProfileRate(rate)
   126  }
   127  
   128  func StartPProf(address string) {
   129  	log.Info("Starting pprof server", "addr", fmt.Sprintf("http://%s/debug/pprof", address))
   130  
   131  	go func() {
   132  		// nolint: gosec
   133  		if err := http.ListenAndServe(address, nil); err != nil {
   134  			log.Error("Failure in running pprof server", "err", err)
   135  		}
   136  	}()
   137  }