github.com/grafana/pyroscope-go/godeltaprof@v0.1.8-0.20240513050943-1b1f97373e2a/heap.go (about)

     1  package godeltaprof
     2  
     3  import (
     4  	"io"
     5  	"runtime"
     6  	"sync"
     7  
     8  	"github.com/grafana/pyroscope-go/godeltaprof/internal/pprof"
     9  )
    10  
    11  // HeapProfiler is a stateful profiler for heap allocations in Go programs.
    12  // It is based on runtime.MemProfile and provides similar functionality to
    13  // pprof.WriteHeapProfile, but with some key differences.
    14  //
    15  // The HeapProfiler tracks the delta of heap allocations since the last
    16  // profile was written, effectively providing a snapshot of the changes
    17  // in heap usage between two points in time. This is in contrast to the
    18  // pprof.WriteHeapProfile function, which accumulates profiling data
    19  // and results in profiles that represent the entire lifetime of the program.
    20  //
    21  // The HeapProfiler is safe for concurrent use, as it serializes access to
    22  // its internal state using a sync.Mutex. This ensures that multiple goroutines
    23  // can call the Profile method without causing any data race issues.
    24  //
    25  // Usage:
    26  //
    27  //	hp := godeltaprof.NewHeapProfiler()
    28  //	...
    29  //	err := hp.Profile(someWriter)
    30  type HeapProfiler struct {
    31  	impl  pprof.DeltaHeapProfiler
    32  	mutex sync.Mutex
    33  }
    34  
    35  func NewHeapProfiler() *HeapProfiler {
    36  	return &HeapProfiler{
    37  		impl: pprof.DeltaHeapProfiler{
    38  			Options: pprof.ProfileBuilderOptions{
    39  				GenericsFrames: true,
    40  				LazyMapping:    true,
    41  			},
    42  		}}
    43  }
    44  
    45  func NewHeapProfilerWithOptions(options ProfileOptions) *HeapProfiler {
    46  	return &HeapProfiler{
    47  		impl: pprof.DeltaHeapProfiler{
    48  			Options: pprof.ProfileBuilderOptions{
    49  				GenericsFrames: options.GenericsFrames,
    50  				LazyMapping:    options.LazyMappings,
    51  			},
    52  		}}
    53  }
    54  
    55  func (d *HeapProfiler) Profile(w io.Writer) error {
    56  	d.mutex.Lock()
    57  	defer d.mutex.Unlock()
    58  
    59  	// Find out how many records there are (MemProfile(nil, true)),
    60  	// allocate that many records, and get the data.
    61  	// There's a race—more records might be added between
    62  	// the two calls—so allocate a few extra records for safety
    63  	// and also try again if we're very unlucky.
    64  	// The loop should only execute one iteration in the common case.
    65  	var p []runtime.MemProfileRecord
    66  	n, ok := runtime.MemProfile(nil, true)
    67  	for {
    68  		// Allocate room for a slightly bigger profile,
    69  		// in case a few more entries have been added
    70  		// since the call to MemProfile.
    71  		p = make([]runtime.MemProfileRecord, n+50)
    72  		n, ok = runtime.MemProfile(p, true)
    73  		if ok {
    74  			p = p[0:n]
    75  			break
    76  		}
    77  		// Profile grew; try again.
    78  	}
    79  
    80  	return d.impl.WriteHeapProto(w, p, int64(runtime.MemProfileRate), "")
    81  }