github.com/bazelbuild/bazel-gazelle@v0.36.1-0.20240520142334-61b277ba6fed/cmd/gazelle/profiler.go (about)

     1  package main
     2  
     3  import (
     4  	"os"
     5  	"runtime"
     6  	"runtime/pprof"
     7  )
     8  
     9  type profiler struct {
    10  	cpuProfile *os.File
    11  	memProfile string
    12  }
    13  
    14  // newProfiler creates a profiler that writes to the given files.
    15  // it returns an empty profiler if both files are empty.
    16  // so that stop() will never fail.
    17  func newProfiler(cpuProfile, memProfile string) (profiler, error) {
    18  	if cpuProfile == "" {
    19  		return profiler{
    20  			memProfile: memProfile,
    21  		}, nil
    22  	}
    23  
    24  	f, err := os.Create(cpuProfile)
    25  	if err != nil {
    26  		return profiler{}, err
    27  	}
    28  	pprof.StartCPUProfile(f)
    29  
    30  	return profiler{
    31  		cpuProfile: f,
    32  		memProfile: memProfile,
    33  	}, nil
    34  }
    35  
    36  func (p *profiler) stop() error {
    37  	if p.cpuProfile != nil {
    38  		pprof.StopCPUProfile()
    39  		p.cpuProfile.Close()
    40  	}
    41  
    42  	if p.memProfile == "" {
    43  		return nil
    44  	}
    45  
    46  	f, err := os.Create(p.memProfile)
    47  	if err != nil {
    48  		return err
    49  	}
    50  	defer f.Close()
    51  	runtime.GC()
    52  	return pprof.WriteHeapProfile(f)
    53  }