github.com/pyroscope-io/godeltaprof@v0.1.3-0.20230906152420-0d7eeca7b8c1/heap.go (about)

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