github.com/fluhus/gostuff@v0.4.1-0.20240331134726-be71864f2b5d/ezpprof/ezpprof.go (about)

     1  // Package ezpprof is a convenience wrapper over the runtime/pprof package.
     2  //
     3  // This package helps to quickly introduce profiling to a piece of code without
     4  // the mess of opening files and checking errors.
     5  //
     6  // A typical use of this package looks like:
     7  //
     8  //	ezpprof.Start("myfile.pprof")
     9  //	{... some complicated code ...}
    10  //	ezpprof.Stop()
    11  //
    12  // Or alternatively:
    13  //
    14  //	const profile = true
    15  //
    16  //	if profile {
    17  //	  ezpprof.Start("myfile.pprof")
    18  //	  defer ezpprof.Stop()
    19  //	}
    20  package ezpprof
    21  
    22  import (
    23  	"io"
    24  	"runtime/pprof"
    25  
    26  	"github.com/fluhus/gostuff/aio"
    27  )
    28  
    29  var fout io.WriteCloser
    30  
    31  // Start starts CPU profiling and writes to the given file.
    32  // Panics if an error occurs.
    33  func Start(file string) {
    34  	if fout != nil {
    35  		panic("already profiling")
    36  	}
    37  	f, err := aio.CreateRaw(file)
    38  	if err != nil {
    39  		panic(err)
    40  	}
    41  	fout = f
    42  	pprof.StartCPUProfile(fout)
    43  }
    44  
    45  // Stop stops CPU profiling and closes the output file.
    46  // Panics if called without calling Start.
    47  func Stop() {
    48  	if fout == nil {
    49  		panic("Stop called without calling Start")
    50  	}
    51  	pprof.StopCPUProfile()
    52  	if err := fout.Close(); err != nil {
    53  		panic(err)
    54  	}
    55  	fout = nil
    56  }
    57  
    58  // Heap writes heap profile to the given file. Panics if an error occurs.
    59  func Heap(file string) {
    60  	f, err := aio.Create(file)
    61  	if err != nil {
    62  		panic(err)
    63  	}
    64  	defer f.Close()
    65  
    66  	err = pprof.WriteHeapProfile(f)
    67  	if err != nil {
    68  		panic(err)
    69  	}
    70  }