github.com/go-asm/go@v1.21.1-0.20240213172139-40c5ead50c48/trace/traceviewer/pprof.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Serving of pprof-like profiles.
     6  
     7  package traceviewer
     8  
     9  import (
    10  	"bufio"
    11  	"fmt"
    12  	"net/http"
    13  	"os"
    14  	"os/exec"
    15  	"path/filepath"
    16  	"runtime"
    17  	"time"
    18  
    19  	"github.com/go-asm/go/profile"
    20  	"github.com/go-asm/go/trace"
    21  )
    22  
    23  type ProfileFunc func(r *http.Request) ([]ProfileRecord, error)
    24  
    25  // SVGProfileHandlerFunc serves pprof-like profile generated by prof as svg.
    26  func SVGProfileHandlerFunc(f ProfileFunc) http.HandlerFunc {
    27  	return func(w http.ResponseWriter, r *http.Request) {
    28  		if r.FormValue("raw") != "" {
    29  			w.Header().Set("Content-Type", "application/octet-stream")
    30  
    31  			failf := func(s string, args ...any) {
    32  				w.Header().Set("Content-Type", "text/plain; charset=utf-8")
    33  				w.Header().Set("X-Go-Pprof", "1")
    34  				http.Error(w, fmt.Sprintf(s, args...), http.StatusInternalServerError)
    35  			}
    36  			records, err := f(r)
    37  			if err != nil {
    38  				failf("failed to get records: %v", err)
    39  				return
    40  			}
    41  			if err := BuildProfile(records).Write(w); err != nil {
    42  				failf("failed to write profile: %v", err)
    43  				return
    44  			}
    45  			return
    46  		}
    47  
    48  		blockf, err := os.CreateTemp("", "block")
    49  		if err != nil {
    50  			http.Error(w, fmt.Sprintf("failed to create temp file: %v", err), http.StatusInternalServerError)
    51  			return
    52  		}
    53  		defer func() {
    54  			blockf.Close()
    55  			os.Remove(blockf.Name())
    56  		}()
    57  		records, err := f(r)
    58  		if err != nil {
    59  			http.Error(w, fmt.Sprintf("failed to generate profile: %v", err), http.StatusInternalServerError)
    60  		}
    61  		blockb := bufio.NewWriter(blockf)
    62  		if err := BuildProfile(records).Write(blockb); err != nil {
    63  			http.Error(w, fmt.Sprintf("failed to write profile: %v", err), http.StatusInternalServerError)
    64  			return
    65  		}
    66  		if err := blockb.Flush(); err != nil {
    67  			http.Error(w, fmt.Sprintf("failed to flush temp file: %v", err), http.StatusInternalServerError)
    68  			return
    69  		}
    70  		if err := blockf.Close(); err != nil {
    71  			http.Error(w, fmt.Sprintf("failed to close temp file: %v", err), http.StatusInternalServerError)
    72  			return
    73  		}
    74  		svgFilename := blockf.Name() + ".svg"
    75  		if output, err := exec.Command(goCmd(), "tool", "pprof", "-svg", "-output", svgFilename, blockf.Name()).CombinedOutput(); err != nil {
    76  			http.Error(w, fmt.Sprintf("failed to execute go tool pprof: %v\n%s", err, output), http.StatusInternalServerError)
    77  			return
    78  		}
    79  		defer os.Remove(svgFilename)
    80  		w.Header().Set("Content-Type", "image/svg+xml")
    81  		http.ServeFile(w, r, svgFilename)
    82  	}
    83  }
    84  
    85  type ProfileRecord struct {
    86  	Stack []*trace.Frame
    87  	Count uint64
    88  	Time  time.Duration
    89  }
    90  
    91  func BuildProfile(prof []ProfileRecord) *profile.Profile {
    92  	p := &profile.Profile{
    93  		PeriodType: &profile.ValueType{Type: "trace", Unit: "count"},
    94  		Period:     1,
    95  		SampleType: []*profile.ValueType{
    96  			{Type: "contentions", Unit: "count"},
    97  			{Type: "delay", Unit: "nanoseconds"},
    98  		},
    99  	}
   100  	locs := make(map[uint64]*profile.Location)
   101  	funcs := make(map[string]*profile.Function)
   102  	for _, rec := range prof {
   103  		var sloc []*profile.Location
   104  		for _, frame := range rec.Stack {
   105  			loc := locs[frame.PC]
   106  			if loc == nil {
   107  				fn := funcs[frame.File+frame.Fn]
   108  				if fn == nil {
   109  					fn = &profile.Function{
   110  						ID:         uint64(len(p.Function) + 1),
   111  						Name:       frame.Fn,
   112  						SystemName: frame.Fn,
   113  						Filename:   frame.File,
   114  					}
   115  					p.Function = append(p.Function, fn)
   116  					funcs[frame.File+frame.Fn] = fn
   117  				}
   118  				loc = &profile.Location{
   119  					ID:      uint64(len(p.Location) + 1),
   120  					Address: frame.PC,
   121  					Line: []profile.Line{
   122  						{
   123  							Function: fn,
   124  							Line:     int64(frame.Line),
   125  						},
   126  					},
   127  				}
   128  				p.Location = append(p.Location, loc)
   129  				locs[frame.PC] = loc
   130  			}
   131  			sloc = append(sloc, loc)
   132  		}
   133  		p.Sample = append(p.Sample, &profile.Sample{
   134  			Value:    []int64{int64(rec.Count), int64(rec.Time)},
   135  			Location: sloc,
   136  		})
   137  	}
   138  	return p
   139  }
   140  
   141  func goCmd() string {
   142  	var exeSuffix string
   143  	if runtime.GOOS == "windows" {
   144  		exeSuffix = ".exe"
   145  	}
   146  	path := filepath.Join(runtime.GOROOT(), "bin", "go"+exeSuffix)
   147  	if _, err := os.Stat(path); err == nil {
   148  		return path
   149  	}
   150  	return "go"
   151  }