github.com/Benchkram/bob@v0.0.0-20220321080157-7c8f3876e225/cli/pprof.go (about)

     1  package cli
     2  
     3  import (
     4  	"os"
     5  	"runtime/pprof"
     6  
     7  	"github.com/Benchkram/errz"
     8  )
     9  
    10  const _cpuprofile = "cpuprofile.prof"
    11  const _memprofile = "memprofile.prof"
    12  
    13  func profiling(cpuprofile, memprofile bool) func() {
    14  	doOnStop := []func(){}
    15  	stop := func() {
    16  		for _, d := range doOnStop {
    17  			if d != nil {
    18  				d()
    19  			}
    20  		}
    21  	}
    22  
    23  	if cpuprofile {
    24  		f, err := os.Create(_cpuprofile)
    25  		errz.Fatal(err)
    26  
    27  		_ = pprof.StartCPUProfile(f)
    28  		doOnStop = append(doOnStop, pprof.StopCPUProfile)
    29  	}
    30  
    31  	if memprofile {
    32  		f, err := os.Create(_memprofile)
    33  		errz.Fatal(err)
    34  
    35  		doOnStop = append(doOnStop, func() {
    36  			_ = pprof.WriteHeapProfile(f)
    37  			_ = f.Close()
    38  		})
    39  	}
    40  
    41  	return stop
    42  }