github.com/shrimpyuk/bor@v0.2.15-0.20220224151350-fb4ec6020bae/internal/cli/server/pprof/pprof.go (about)

     1  package pprof
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"runtime"
     8  	"runtime/pprof"
     9  	"runtime/trace"
    10  	"time"
    11  )
    12  
    13  // Profile generates a pprof.Profile report for the given profile name.
    14  func Profile(profile string, debug, gc int) ([]byte, map[string]string, error) {
    15  	p := pprof.Lookup(profile)
    16  	if p == nil {
    17  		return nil, nil, fmt.Errorf("profile '%s' not found", profile)
    18  	}
    19  
    20  	if profile == "heap" && gc > 0 {
    21  		runtime.GC()
    22  	}
    23  
    24  	var buf bytes.Buffer
    25  	if err := p.WriteTo(&buf, debug); err != nil {
    26  		return nil, nil, err
    27  	}
    28  
    29  	headers := map[string]string{
    30  		"X-Content-Type-Options": "nosniff",
    31  	}
    32  	if debug != 0 {
    33  		headers["Content-Type"] = "text/plain; charset=utf-8"
    34  	} else {
    35  		headers["Content-Type"] = "application/octet-stream"
    36  		headers["Content-Disposition"] = fmt.Sprintf(`attachment; filename="%s"`, profile)
    37  	}
    38  	return buf.Bytes(), headers, nil
    39  }
    40  
    41  // CPUProfile generates a CPU Profile for a given duration
    42  func CPUProfile(ctx context.Context, sec int) ([]byte, map[string]string, error) {
    43  	if sec <= 0 {
    44  		sec = 1
    45  	}
    46  
    47  	var buf bytes.Buffer
    48  	if err := pprof.StartCPUProfile(&buf); err != nil {
    49  		return nil, nil, err
    50  	}
    51  
    52  	sleep(ctx, time.Duration(sec)*time.Second)
    53  
    54  	pprof.StopCPUProfile()
    55  
    56  	return buf.Bytes(),
    57  		map[string]string{
    58  			"X-Content-Type-Options": "nosniff",
    59  			"Content-Type":           "application/octet-stream",
    60  			"Content-Disposition":    `attachment; filename="profile"`,
    61  		}, nil
    62  }
    63  
    64  // Trace runs a trace profile for a given duration
    65  func Trace(ctx context.Context, sec int) ([]byte, map[string]string, error) {
    66  	if sec <= 0 {
    67  		sec = 1
    68  	}
    69  
    70  	var buf bytes.Buffer
    71  	if err := trace.Start(&buf); err != nil {
    72  		return nil, nil, err
    73  	}
    74  
    75  	sleep(ctx, time.Duration(sec)*time.Second)
    76  
    77  	trace.Stop()
    78  
    79  	return buf.Bytes(),
    80  		map[string]string{
    81  			"X-Content-Type-Options": "nosniff",
    82  			"Content-Type":           "application/octet-stream",
    83  			"Content-Disposition":    `attachment; filename="trace"`,
    84  		}, nil
    85  }
    86  
    87  func sleep(ctx context.Context, d time.Duration) {
    88  	// Sleep until duration is met or ctx is cancelled
    89  	select {
    90  	case <-time.After(d):
    91  	case <-ctx.Done():
    92  	}
    93  }