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

     1  package wzprof
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"net/http"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/google/pprof/profile"
    11  	"github.com/tetratelabs/wazero/api"
    12  	"github.com/tetratelabs/wazero/experimental"
    13  )
    14  
    15  // MemoryProfiler is the implementation of a performance profiler recording
    16  // samples of memory allocation and utilization.
    17  //
    18  // The profiler generates the following samples:
    19  // - "alloc_objects" records the locations where objects are allocated
    20  // - "alloc_space"   records the locations where bytes are allocated
    21  // - "inuse_objects" records the allocation of active objects
    22  // - "inuse_space"   records the bytes used by active objects
    23  //
    24  // "alloc_objects" and "alloc_space" are all time counters since the start of
    25  // the program, while "inuse_objects" and "inuse_space" capture the current state
    26  // of the program at the time the profile is taken.
    27  type MemoryProfiler struct {
    28  	p     *Profiling
    29  	mutex sync.Mutex
    30  	alloc stackCounterMap
    31  	inuse map[uint32]memoryAllocation
    32  	start time.Time
    33  }
    34  
    35  // MemoryProfilerOption is a type used to represent configuration options for
    36  // MemoryProfiler instances created by NewMemoryProfiler.
    37  type MemoryProfilerOption func(*MemoryProfiler)
    38  
    39  // InuseMemory is a memory profiler option which enables tracking of allocated
    40  // and freed objects to generate snapshots of the current state of a program
    41  // memory.
    42  func InuseMemory(enable bool) MemoryProfilerOption {
    43  	return func(p *MemoryProfiler) {
    44  		if enable {
    45  			p.inuse = make(map[uint32]memoryAllocation)
    46  		}
    47  	}
    48  }
    49  
    50  type memoryAllocation struct {
    51  	*stackCounter
    52  	size uint32
    53  }
    54  
    55  // newMemoryProfiler constructs a new instance of MemoryProfiler using the given
    56  // time function to record the profile execution time.
    57  func newMemoryProfiler(p *Profiling, options ...MemoryProfilerOption) *MemoryProfiler {
    58  	m := &MemoryProfiler{
    59  		p:     p,
    60  		alloc: make(stackCounterMap),
    61  		start: time.Now(),
    62  	}
    63  	for _, opt := range options {
    64  		opt(m)
    65  	}
    66  	return m
    67  }
    68  
    69  // NewProfile takes a snapshot of the current memory allocation state and builds
    70  // a profile representing the state of the program memory.
    71  func (p *MemoryProfiler) NewProfile(sampleRate float64) *profile.Profile {
    72  	ratio := 1 / sampleRate
    73  	return buildProfile(p.p, p.snapshot(), p.start, time.Since(p.start), p.SampleType(),
    74  		[]float64{ratio, ratio, ratio, ratio},
    75  	)
    76  }
    77  
    78  // Name returns "allocs" to match the name of the memory profiler in pprof.
    79  func (p *MemoryProfiler) Name() string {
    80  	return "allocs"
    81  }
    82  
    83  // Desc returns a description copied from net/http/pprof.
    84  func (p *MemoryProfiler) Desc() string {
    85  	return profileDescriptions[p.Name()]
    86  }
    87  
    88  // Count returns the number of allocation stacks recorded in p.
    89  func (p *MemoryProfiler) Count() int {
    90  	p.mutex.Lock()
    91  	n := p.alloc.len()
    92  	p.mutex.Unlock()
    93  	return n
    94  }
    95  
    96  // SampleType returns the set of value types present in samples recorded by the
    97  // memory profiler.
    98  func (p *MemoryProfiler) SampleType() []*profile.ValueType {
    99  	sampleType := []*profile.ValueType{
   100  		{Type: "alloc_objects", Unit: "count"},
   101  		{Type: "alloc_space", Unit: "bytes"},
   102  	}
   103  
   104  	if p.inuse != nil {
   105  		// TODO: when can track freeing of garbage collected languages like Go,
   106  		// this should be enabled by default, and we can remove the slicing of
   107  		// sample values in buildProfile.
   108  		sampleType = append(sampleType,
   109  			&profile.ValueType{Type: "inuse_objects", Unit: "count"},
   110  			&profile.ValueType{Type: "inuse_space", Unit: "bytes"},
   111  		)
   112  	}
   113  
   114  	return sampleType
   115  }
   116  
   117  type memorySample struct {
   118  	stack stackTrace
   119  	value [4]int64 // allocCount, allocBytes, inuseCount, inuseBytes
   120  }
   121  
   122  func (m *memorySample) sampleLocation() stackTrace {
   123  	return m.stack
   124  }
   125  
   126  func (m *memorySample) sampleValue() []int64 {
   127  	return m.value[:]
   128  }
   129  
   130  func (p *MemoryProfiler) snapshot() map[uint64]*memorySample {
   131  	// We hold an exclusive lock while getting a snapshot of the profiler state.
   132  	// This will block concurrent calls to malloc/free/etc... We accept the cost
   133  	// since it only happens when the memory profile is captured, and memory
   134  	// allocation is generally accepted as being a potentially costly operation.
   135  	p.mutex.Lock()
   136  	defer p.mutex.Unlock()
   137  
   138  	samples := make(map[uint64]*memorySample, len(p.alloc))
   139  
   140  	for _, alloc := range p.alloc {
   141  		p := samples[alloc.stack.key]
   142  		if p == nil {
   143  			p = &memorySample{stack: alloc.stack}
   144  			samples[alloc.stack.key] = p
   145  		}
   146  		p.value[0] += alloc.count()
   147  		p.value[1] += alloc.total()
   148  	}
   149  
   150  	for _, inuse := range p.inuse {
   151  		p := samples[inuse.stack.key]
   152  		p.value[2] += 1
   153  		p.value[3] += int64(inuse.size)
   154  	}
   155  
   156  	return samples
   157  }
   158  
   159  // NewHandler returns a http handler allowing the profiler to be exposed on a
   160  // pprof-compatible http endpoint.
   161  //
   162  // The sample rate is a value between 0 and 1 used to scale the profile results
   163  // based on the sampling rate applied to the profiler so the resulting values
   164  // remain representative.
   165  //
   166  // The symbolizer passed as argument is used to resolve names of program
   167  // locations recorded in the profile.
   168  func (p *MemoryProfiler) NewHandler(sampleRate float64) http.Handler {
   169  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   170  		serveProfile(w, p.NewProfile(sampleRate))
   171  	})
   172  }
   173  
   174  // NewFunctionListener returns a function listener suited to install a hook on
   175  // functions responsible for memory allocation.
   176  //
   177  // The listener recognizes multiple memory allocation functions used by
   178  // compilers and libraries. It uses the function name to detect memory
   179  // allocators, currently supporting libc, Go, and TinyGo.
   180  func (p *MemoryProfiler) NewFunctionListener(def api.FunctionDefinition) experimental.FunctionListener {
   181  	if p.p.lang == python311 {
   182  		switch def.Name() {
   183  		// Raw domain
   184  		case "PyMem_RawMalloc":
   185  			return profilingListener{p.p, &mallocProfiler{memory: p}}
   186  		case "PyMem_RawCalloc":
   187  			return profilingListener{p.p, &callocProfiler{memory: p}}
   188  		case "PyMem_RawRealloc":
   189  			return profilingListener{p.p, &reallocProfiler{memory: p}}
   190  		case "PyMem_RawFree":
   191  			return profilingListener{p.p, &freeProfiler{memory: p}}
   192  		// Memory domain
   193  		case "PyMem_Malloc":
   194  			return profilingListener{p.p, &mallocProfiler{memory: p}}
   195  		case "PyMem_Calloc":
   196  			return profilingListener{p.p, &callocProfiler{memory: p}}
   197  		case "PyMem_Realloc":
   198  			return profilingListener{p.p, &reallocProfiler{memory: p}}
   199  		case "PyMem_Free":
   200  			return profilingListener{p.p, &freeProfiler{memory: p}}
   201  		// Object domain
   202  		case "PyObject_Malloc":
   203  			return profilingListener{p.p, &mallocProfiler{memory: p}}
   204  		case "PyObject_Calloc":
   205  			return profilingListener{p.p, &callocProfiler{memory: p}}
   206  		case "PyObject_Realloc":
   207  			return profilingListener{p.p, &reallocProfiler{memory: p}}
   208  		case "PyObject_Free":
   209  			return profilingListener{p.p, &freeProfiler{memory: p}}
   210  		}
   211  		return nil
   212  	}
   213  	switch def.Name() {
   214  	// C standard library, Rust
   215  	case "malloc":
   216  		return profilingListener{p.p, &mallocProfiler{memory: p}}
   217  	case "calloc":
   218  		return profilingListener{p.p, &callocProfiler{memory: p}}
   219  	case "realloc":
   220  		return profilingListener{p.p, &reallocProfiler{memory: p}}
   221  	case "free":
   222  		return profilingListener{p.p, &freeProfiler{memory: p}}
   223  
   224  	// Go
   225  	case "runtime.mallocgc":
   226  		return profilingListener{p.p, &goRuntimeMallocgcProfiler{memory: p}}
   227  
   228  	// TinyGo
   229  	case "runtime.alloc":
   230  		return profilingListener{p.p, &mallocProfiler{memory: p}}
   231  
   232  	default:
   233  		return nil
   234  	}
   235  }
   236  
   237  func (p *MemoryProfiler) observeAlloc(addr, size uint32, stack stackTrace) {
   238  	p.mutex.Lock()
   239  	alloc := p.alloc.lookup(stack)
   240  	alloc.observe(int64(size))
   241  	if p.inuse != nil {
   242  		p.inuse[addr] = memoryAllocation{alloc, size}
   243  	}
   244  	p.mutex.Unlock()
   245  }
   246  
   247  func (p *MemoryProfiler) observeFree(addr uint32) {
   248  	if p.inuse != nil {
   249  		p.mutex.Lock()
   250  		delete(p.inuse, addr)
   251  		p.mutex.Unlock()
   252  	}
   253  }
   254  
   255  type mallocProfiler struct {
   256  	memory *MemoryProfiler
   257  	size   uint32
   258  	stack  stackTrace
   259  }
   260  
   261  func (p *mallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
   262  	p.size = api.DecodeU32(params[0])
   263  	p.stack = makeStackTrace(p.stack, si)
   264  }
   265  
   266  func (p *mallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
   267  	p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
   268  }
   269  
   270  func (p *mallocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
   271  }
   272  
   273  type callocProfiler struct {
   274  	memory *MemoryProfiler
   275  	count  uint32
   276  	size   uint32
   277  	stack  stackTrace
   278  }
   279  
   280  func (p *callocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
   281  	p.count = api.DecodeU32(params[0])
   282  	p.size = api.DecodeU32(params[1])
   283  	p.stack = makeStackTrace(p.stack, si)
   284  }
   285  
   286  func (p *callocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
   287  	p.memory.observeAlloc(api.DecodeU32(results[0]), p.count*p.size, p.stack)
   288  }
   289  
   290  func (p *callocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
   291  }
   292  
   293  type reallocProfiler struct {
   294  	memory *MemoryProfiler
   295  	addr   uint32
   296  	size   uint32
   297  	stack  stackTrace
   298  }
   299  
   300  func (p *reallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
   301  	p.addr = api.DecodeU32(params[0])
   302  	p.size = api.DecodeU32(params[1])
   303  	p.stack = makeStackTrace(p.stack, si)
   304  }
   305  
   306  func (p *reallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
   307  	p.memory.observeFree(p.addr)
   308  	p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
   309  }
   310  
   311  func (p *reallocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
   312  }
   313  
   314  type freeProfiler struct {
   315  	memory *MemoryProfiler
   316  	addr   uint32
   317  }
   318  
   319  func (p *freeProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
   320  	p.addr = api.DecodeU32(params[0])
   321  }
   322  
   323  func (p *freeProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64) {
   324  	p.memory.observeFree(p.addr)
   325  }
   326  
   327  func (p *freeProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
   328  	p.After(ctx, mod, def, nil)
   329  }
   330  
   331  type goRuntimeMallocgcProfiler struct {
   332  	memory *MemoryProfiler
   333  	size   uint32
   334  	stack  stackTrace
   335  }
   336  
   337  func (p *goRuntimeMallocgcProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, wasmsi experimental.StackIterator) {
   338  	imod := mod.(experimental.InternalModule)
   339  	mem := imod.Memory()
   340  
   341  	sp := uint32(imod.Global(0).Get())
   342  	offset := sp + 8*(uint32(0)+1) // +1 for the return address
   343  	b, ok := mem.Read(offset, 8)
   344  	if ok {
   345  		p.size = binary.LittleEndian.Uint32(b)
   346  		p.stack = makeStackTrace(p.stack, wasmsi)
   347  	} else {
   348  		p.size = 0
   349  	}
   350  }
   351  
   352  func (p *goRuntimeMallocgcProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64) {
   353  	if p.size != 0 {
   354  		// TODO: get the returned pointer
   355  		addr := uint32(0)
   356  		p.memory.observeAlloc(addr, p.size, p.stack)
   357  	}
   358  }
   359  
   360  func (p *goRuntimeMallocgcProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
   361  	p.After(ctx, mod, def, nil)
   362  }