src.elv.sh@v0.21.0-dev.0.20240515223629-06979efb9a2a/pkg/pprof/pprof.go (about)

     1  // Package pprof adds profiling support to the Elvish program.
     2  package pprof
     3  
     4  import (
     5  	"fmt"
     6  	"os"
     7  	"runtime/pprof"
     8  
     9  	"src.elv.sh/pkg/prog"
    10  )
    11  
    12  // Program adds support for the -cpuprofile flag.
    13  type Program struct {
    14  	cpuProfile    string
    15  	allocsProfile string
    16  }
    17  
    18  func (p *Program) RegisterFlags(f *prog.FlagSet) {
    19  	f.StringVar(&p.cpuProfile, "cpuprofile", "", "write CPU profile to file")
    20  	f.StringVar(&p.allocsProfile, "allocsprofile", "", "write memory allocation profile to file")
    21  }
    22  
    23  func (p *Program) Run(fds [3]*os.File, _ []string) error {
    24  	var cleanups []func([3]*os.File)
    25  	if p.cpuProfile != "" {
    26  		f, err := os.Create(p.cpuProfile)
    27  		if err != nil {
    28  			fmt.Fprintln(fds[2], "Warning: cannot create CPU profile:", err)
    29  			fmt.Fprintln(fds[2], "Continuing without CPU profiling.")
    30  		} else {
    31  			pprof.StartCPUProfile(f)
    32  			cleanups = append(cleanups, func([3]*os.File) {
    33  				pprof.StopCPUProfile()
    34  				f.Close()
    35  			})
    36  		}
    37  	}
    38  	if p.allocsProfile != "" {
    39  		f, err := os.Create(p.allocsProfile)
    40  		if err != nil {
    41  			fmt.Fprintln(fds[2], "Warning: cannot create memory allocation profile:", err)
    42  			fmt.Fprintln(fds[2], "Continuing without memory allocation profiling.")
    43  		} else {
    44  			cleanups = append(cleanups, func([3]*os.File) {
    45  				pprof.Lookup("allocs").WriteTo(f, 0)
    46  				f.Close()
    47  			})
    48  		}
    49  	}
    50  	return prog.NextProgram(cleanups...)
    51  }
    52  
    53  func (p *Program) Cleanup() {
    54  	if p.cpuProfile != "" {
    55  		pprof.StopCPUProfile()
    56  	}
    57  }