github.com/stealthrocket/wzprof@v0.2.1-0.20230830205924-5fa86be5e5b3/wzprof.go (about)

     1  package wzprof
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"hash/maphash"
     7  	"net/http"
     8  	"os"
     9  	"strings"
    10  	"time"
    11  	"unsafe"
    12  
    13  	"github.com/google/pprof/profile"
    14  	"github.com/tetratelabs/wazero"
    15  	"github.com/tetratelabs/wazero/api"
    16  	"github.com/tetratelabs/wazero/experimental"
    17  	"golang.org/x/exp/slices"
    18  )
    19  
    20  // Profiling mechanism for a given WASM binary. Entry point to generate
    21  // Profilers.
    22  type Profiling struct {
    23  	wasm []byte
    24  
    25  	onlyFunctions     map[string]struct{}
    26  	filteredFunctions map[string]struct{}
    27  	symbols           symbolizer
    28  	stackIterator     func(mod api.Module, def api.FunctionDefinition, wasmsi experimental.StackIterator) experimental.StackIterator
    29  
    30  	lang language
    31  }
    32  
    33  type language int8
    34  
    35  const (
    36  	unknown language = iota
    37  	golang
    38  	python311
    39  )
    40  
    41  // ProfilingFor a given wasm binary. The resulting Profiling needs to be
    42  // prepared after Wazero module compilation.
    43  func ProfilingFor(wasm []byte) *Profiling {
    44  	r := &Profiling{
    45  		wasm:    wasm,
    46  		symbols: noopsymbolizer{},
    47  		stackIterator: func(mod api.Module, def api.FunctionDefinition, wasmsi experimental.StackIterator) experimental.StackIterator {
    48  			return wasmsi
    49  		},
    50  	}
    51  
    52  	if binCompiledByGo(wasm) {
    53  		r.lang = golang
    54  		// Those functions are special. They use a different calling
    55  		// convention. Their call sites do not update the stack pointer,
    56  		// which makes it impossible to correctly walk the stack.
    57  		//
    58  		// https://github.com/golang/go/blob/7ad92e95b56019083824492fbec5bb07926d8ebd/src/cmd/internal/obj/wasm/wasmobj.go#LL907C18-L930C2
    59  		r.filteredFunctions = map[string]struct{}{
    60  			"_rt0_wasm_js":            {},
    61  			"_rt0_wasm_wasip1":        {},
    62  			"wasm_export_run":         {},
    63  			"wasm_export_resume":      {},
    64  			"wasm_export_getsp":       {},
    65  			"wasm_pc_f_loop":          {},
    66  			"gcWriteBarrier":          {},
    67  			"runtime.gcWriteBarrier1": {},
    68  			"runtime.gcWriteBarrier2": {},
    69  			"runtime.gcWriteBarrier3": {},
    70  			"runtime.gcWriteBarrier4": {},
    71  			"runtime.gcWriteBarrier5": {},
    72  			"runtime.gcWriteBarrier6": {},
    73  			"runtime.gcWriteBarrier7": {},
    74  			"runtime.gcWriteBarrier8": {},
    75  			"runtime.wasmDiv":         {},
    76  			"runtime.wasmTruncS":      {},
    77  			"runtime.wasmTruncU":      {},
    78  			"cmpbody":                 {},
    79  			"memeqbody":               {},
    80  			"memcmp":                  {},
    81  			"memchr":                  {},
    82  		}
    83  	} else if supportedPython(wasm) {
    84  		r.lang = python311
    85  		r.onlyFunctions = map[string]struct{}{
    86  			"PyObject_Vectorcall": {},
    87  			// Those functions are also likely candidate for useful profiling.
    88  			// We may need to look into them if someone reports missing frames.
    89  			//
    90  			// "_PyEval_EvalFrameDefault": {},
    91  			// "_PyEvalFramePushAndInit": {},
    92  		}
    93  	}
    94  
    95  	return r
    96  }
    97  
    98  // CPUProfiler constructs a new instance of CPUProfiler using the given time
    99  // function to record the CPU time consumed.
   100  func (p *Profiling) CPUProfiler(options ...CPUProfilerOption) *CPUProfiler {
   101  	return newCPUProfiler(p, options...)
   102  }
   103  
   104  // MemoryProfiler constructs a new instance of MemoryProfiler using the given
   105  // time function to record the profile execution time.
   106  func (p *Profiling) MemoryProfiler(options ...MemoryProfilerOption) *MemoryProfiler {
   107  	return newMemoryProfiler(p, options...)
   108  }
   109  
   110  // Prepare selects the most appropriate analysis functions for the guest
   111  // code in the provided module.
   112  func (p *Profiling) Prepare(mod wazero.CompiledModule) error {
   113  	switch p.lang {
   114  	case golang:
   115  		s, err := preparePclntabSymbolizer(p.wasm, mod)
   116  		if err != nil {
   117  			return err
   118  		}
   119  
   120  		p.symbols = s
   121  		si := &goStackIterator{
   122  			pclntab:  s,
   123  			unwinder: unwinder{symbols: s},
   124  		}
   125  		p.stackIterator = func(mod api.Module, def api.FunctionDefinition, wasmsi experimental.StackIterator) experimental.StackIterator {
   126  			imod := mod.(experimental.InternalModule)
   127  			si.mem = imod.Memory()
   128  			si.pclntab.EnsureReady(si.mem)
   129  			sp0 := uint32(imod.Global(0).Get())
   130  			gp0 := imod.Global(2).Get()
   131  			pc0 := si.symbols.FIDToPC(fid(def.Index()))
   132  			si.initAt(ptr64(pc0), ptr64(sp0), 0, gptr(gp0), 0)
   133  			si.first = true
   134  			return si
   135  		}
   136  	case python311:
   137  		py, err := preparePython(mod)
   138  		if err != nil {
   139  			return err
   140  		}
   141  		p.symbols = py
   142  		p.stackIterator = py.Stackiter
   143  	default:
   144  		dwarf, err := newDwarfparser(mod)
   145  		if err != nil {
   146  			return nil // TODO: surface error as warning?
   147  		}
   148  		p.symbols = buildDwarfSymbolizer(dwarf)
   149  	}
   150  	return nil
   151  }
   152  
   153  // profilingListener wraps a FunctionListener to adapt its stack iterator to the
   154  // appropriate implementation according to the module support.
   155  type profilingListener struct {
   156  	s *Profiling
   157  	l experimental.FunctionListener
   158  }
   159  
   160  func (s profilingListener) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
   161  	si = s.s.stackIterator(mod, def, si)
   162  	s.l.Before(ctx, mod, def, params, si)
   163  }
   164  
   165  func (s profilingListener) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
   166  	s.l.After(ctx, mod, def, results)
   167  }
   168  
   169  func (s profilingListener) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error) {
   170  	s.l.Abort(ctx, mod, def, err)
   171  }
   172  
   173  // Profiler is an interface implemented by all profiler types available in this
   174  // package.
   175  type Profiler interface {
   176  	experimental.FunctionListenerFactory
   177  
   178  	// Name of the profiler.
   179  	Name() string
   180  
   181  	// Desc is a human-readable description of the profiler.
   182  	Desc() string
   183  
   184  	// Count the number of execution stacks recorded in the profiler.
   185  	Count() int
   186  
   187  	// SampleType returns the set of value types present in samples recorded by
   188  	// the profiler.
   189  	SampleType() []*profile.ValueType
   190  
   191  	// NewHandler returns a new http handler suited to expose profiles on a
   192  	// pprof endpoint.
   193  	NewHandler(sampleRate float64) http.Handler
   194  }
   195  
   196  var (
   197  	_ Profiler = (*CPUProfiler)(nil)
   198  	_ Profiler = (*MemoryProfiler)(nil)
   199  )
   200  
   201  //go:linkname nanotime runtime.nanotime
   202  func nanotime() int64
   203  
   204  // WriteProfile writes a profile to a file at the given path.
   205  func WriteProfile(path string, prof *profile.Profile) error {
   206  	w, err := os.Create(path)
   207  	if err != nil {
   208  		return err
   209  	}
   210  	defer w.Close()
   211  	return prof.Write(w)
   212  }
   213  
   214  type symbolizer interface {
   215  	// Locations returns a list of function locations for a given program
   216  	// counter, and the address it found them at. Locations start from
   217  	// current function followed by the inlined functions, in order of
   218  	// inlining. Result if empty if the pc cannot be resolved.
   219  	Locations(fn experimental.InternalFunction, pc experimental.ProgramCounter) (uint64, []location)
   220  }
   221  
   222  type noopsymbolizer struct{}
   223  
   224  func (s noopsymbolizer) Locations(fn experimental.InternalFunction, pc experimental.ProgramCounter) (uint64, []location) {
   225  	return 0, nil
   226  }
   227  
   228  type location struct {
   229  	File    string
   230  	Line    int64
   231  	Column  int64
   232  	Inlined bool
   233  	// Linkage Name if present, Name otherwise.
   234  	// Only present for inlined functions.
   235  	StableName string
   236  	HumanName  string
   237  }
   238  
   239  func locationForCall(p *Profiling, fn experimental.InternalFunction, pc experimental.ProgramCounter, funcs map[string]*profile.Function) *profile.Location {
   240  	// Cache miss. Get or create function and all the line
   241  	// locations associated with inlining.
   242  	var locations []location
   243  	var symbolFound bool
   244  	def := fn.Definition()
   245  
   246  	out := &profile.Location{}
   247  
   248  	if pc > 0 {
   249  		out.Address, locations = p.symbols.Locations(fn, pc)
   250  		symbolFound = len(locations) > 0
   251  	}
   252  	if len(locations) == 0 {
   253  		// If we don't have a source location, attach to a
   254  		// generic location within the function.
   255  		locations = []location{{}}
   256  	}
   257  	// Provide defaults in case we couldn't resolve DWARF information for
   258  	// the main function call's PC.
   259  	if locations[0].StableName == "" {
   260  		locations[0].StableName = def.Name()
   261  	}
   262  	if locations[0].HumanName == "" {
   263  		locations[0].HumanName = def.Name()
   264  	}
   265  
   266  	lines := make([]profile.Line, len(locations))
   267  
   268  	for i, loc := range locations {
   269  		pprofFn := funcs[loc.StableName]
   270  
   271  		if pprofFn == nil {
   272  			pprofFn = &profile.Function{
   273  				ID:         uint64(len(funcs)) + 1, // 0 is reserved by pprof
   274  				Name:       loc.HumanName,
   275  				SystemName: loc.StableName,
   276  				Filename:   loc.File,
   277  			}
   278  			funcs[loc.StableName] = pprofFn
   279  		} else if symbolFound {
   280  			// Sometimes the function had to be created while the PC
   281  			// wasn't found by the symbol mapper. Attempt to correct
   282  			// it if we had a successful match this time.
   283  			pprofFn.Name = locations[i].HumanName
   284  			pprofFn.SystemName = locations[i].StableName
   285  			pprofFn.Filename = locations[i].File
   286  		}
   287  
   288  		// Pprof expects lines to start with the root of the inlined
   289  		// calls. DWARF encodes that information the other way around,
   290  		// so we fill lines backwards.
   291  		lines[len(locations)-(i+1)] = profile.Line{
   292  			Function: pprofFn,
   293  			Line:     loc.Line,
   294  		}
   295  	}
   296  
   297  	out.Line = lines
   298  	return out
   299  }
   300  
   301  type locationKey struct {
   302  	module string
   303  	index  uint32
   304  	name   string
   305  	pc     uint64
   306  }
   307  
   308  func makeLocationKey(fn api.FunctionDefinition, pc experimental.ProgramCounter) locationKey {
   309  	return locationKey{
   310  		module: fn.ModuleName(),
   311  		index:  fn.Index(),
   312  		name:   fn.Name(),
   313  		pc:     uint64(pc),
   314  	}
   315  }
   316  
   317  type stackCounterMap map[uint64]*stackCounter
   318  
   319  func (scm stackCounterMap) lookup(st stackTrace) *stackCounter {
   320  	sc := scm[st.key]
   321  	if sc == nil {
   322  		sc = &stackCounter{stack: st.clone()}
   323  		scm[st.key] = sc
   324  	}
   325  	return sc
   326  }
   327  
   328  func (scm stackCounterMap) observe(st stackTrace, val int64) {
   329  	scm.lookup(st).observe(val)
   330  }
   331  
   332  func (scm stackCounterMap) len() int {
   333  	return len(scm)
   334  }
   335  
   336  type stackCounter struct {
   337  	stack stackTrace
   338  	value [2]int64 // count, total
   339  }
   340  
   341  func (sc *stackCounter) observe(value int64) {
   342  	sc.value[0] += 1
   343  	sc.value[1] += value
   344  }
   345  
   346  func (sc *stackCounter) count() int64 {
   347  	return sc.value[0]
   348  }
   349  
   350  func (sc *stackCounter) total() int64 {
   351  	return sc.value[1]
   352  }
   353  
   354  func (sc *stackCounter) sampleLocation() stackTrace {
   355  	return sc.stack
   356  }
   357  
   358  func (sc *stackCounter) sampleValue() []int64 {
   359  	return sc.value[:]
   360  }
   361  
   362  func (sc *stackCounter) String() string {
   363  	return fmt.Sprintf("{count:%d,total:%d}", sc.count(), sc.total())
   364  }
   365  
   366  // Compile-time check that program counters are uint64 values.
   367  var _ = assertTypeIsUint64[experimental.ProgramCounter]()
   368  
   369  func assertTypeIsUint64[T ~uint64]() bool {
   370  	return true
   371  }
   372  
   373  type stackFrame struct {
   374  	fn experimental.InternalFunction
   375  	pc experimental.ProgramCounter
   376  }
   377  
   378  type stackTrace struct {
   379  	fns []experimental.InternalFunction
   380  	pcs []experimental.ProgramCounter
   381  	key uint64
   382  }
   383  
   384  func makeStackTrace(st stackTrace, si experimental.StackIterator) stackTrace {
   385  	st.fns = st.fns[:0]
   386  	st.pcs = st.pcs[:0]
   387  
   388  	for si.Next() {
   389  		st.fns = append(st.fns, si.Function())
   390  		st.pcs = append(st.pcs, si.ProgramCounter())
   391  	}
   392  	st.key = maphash.Bytes(stackTraceHashSeed, st.bytes())
   393  	return st
   394  }
   395  
   396  func (st stackTrace) host() bool {
   397  	return len(st.fns) > 0 && st.fns[0].Definition().GoFunction() != nil
   398  }
   399  
   400  func (st stackTrace) len() int {
   401  	return len(st.pcs)
   402  }
   403  
   404  func (st stackTrace) index(i int) stackFrame {
   405  	return stackFrame{
   406  		fn: st.fns[i],
   407  		pc: st.pcs[i],
   408  	}
   409  }
   410  
   411  func (st stackTrace) clone() stackTrace {
   412  	return stackTrace{
   413  		fns: slices.Clone(st.fns),
   414  		pcs: slices.Clone(st.pcs),
   415  		key: st.key,
   416  	}
   417  }
   418  
   419  func (st stackTrace) bytes() []byte {
   420  	pcs := unsafe.SliceData(st.pcs)
   421  	return unsafe.Slice((*byte)(unsafe.Pointer(pcs)), 8*len(st.pcs))
   422  }
   423  
   424  func (st stackTrace) String() string {
   425  	sb := new(strings.Builder)
   426  	for i, n := 0, st.len(); i < n; i++ {
   427  		frame := st.index(i)
   428  		fndef := frame.fn.Definition()
   429  		fmt.Fprintf(sb, "%016x: %s\n", frame.pc, fndef.DebugName())
   430  	}
   431  	return sb.String()
   432  }
   433  
   434  var stackTraceHashSeed = maphash.MakeSeed()
   435  
   436  type sampleType interface {
   437  	sampleLocation() stackTrace
   438  	sampleValue() []int64
   439  }
   440  
   441  func buildProfile[T sampleType](p *Profiling, samples map[uint64]T, start time.Time, duration time.Duration, sampleType []*profile.ValueType, ratios []float64) *profile.Profile {
   442  	prof := &profile.Profile{
   443  		SampleType:    sampleType,
   444  		Sample:        make([]*profile.Sample, 0, len(samples)),
   445  		TimeNanos:     start.UnixNano(),
   446  		DurationNanos: int64(duration),
   447  	}
   448  
   449  	locationID := uint64(1)
   450  	locationCache := make(map[locationKey]*profile.Location)
   451  	functionCache := make(map[string]*profile.Function)
   452  
   453  	for _, sample := range samples {
   454  		stack := sample.sampleLocation()
   455  		location := make([]*profile.Location, stack.len())
   456  
   457  		for i := range location {
   458  			fn := stack.fns[i]
   459  			pc := stack.pcs[i]
   460  
   461  			def := fn.Definition()
   462  			key := makeLocationKey(def, pc)
   463  			loc := locationCache[key]
   464  			if loc == nil {
   465  				loc = locationForCall(p, fn, pc, functionCache)
   466  				loc.ID = locationID
   467  				locationID++
   468  				locationCache[key] = loc
   469  			}
   470  
   471  			location[i] = loc
   472  		}
   473  
   474  		prof.Sample = append(prof.Sample, &profile.Sample{
   475  			Location: location,
   476  			Value:    sample.sampleValue()[:len(sampleType)],
   477  		})
   478  	}
   479  
   480  	prof.Location = make([]*profile.Location, len(locationCache))
   481  	prof.Function = make([]*profile.Function, len(functionCache))
   482  
   483  	for _, loc := range locationCache {
   484  		prof.Location[loc.ID-1] = loc
   485  	}
   486  
   487  	for _, fn := range functionCache {
   488  		prof.Function[fn.ID-1] = fn
   489  	}
   490  
   491  	if err := prof.ScaleN(ratios[:len(sampleType)]); err != nil {
   492  		panic(err)
   493  	}
   494  	return prof
   495  }