github.com/ndau/noms@v1.0.5/go/util/profile/profile.go (about)

     1  // Copyright 2016 Attic Labs, Inc. All rights reserved.
     2  // Licensed under the Apache License, version 2.0:
     3  // http://www.apache.org/licenses/LICENSE-2.0
     4  
     5  package profile
     6  
     7  import (
     8  	"io"
     9  	"os"
    10  	"runtime"
    11  	"runtime/pprof"
    12  
    13  	"github.com/attic-labs/kingpin"
    14  
    15  	"github.com/ndau/noms/go/d"
    16  )
    17  
    18  var (
    19  	cpuProfile   string
    20  	memProfile   string
    21  	blockProfile string
    22  )
    23  
    24  func RegisterProfileFlags(app *kingpin.Application) {
    25  	// Must reset globals because under test this can get called multiple times.
    26  	cpuProfile = ""
    27  	memProfile = ""
    28  	blockProfile = ""
    29  	app.Flag("cpuprofile", "write cpu profile to file").StringVar(&cpuProfile)
    30  	app.Flag("memprofile", "write memory profile to file").StringVar(&memProfile)
    31  	app.Flag("blockprofile", "write block profile to file").StringVar(&blockProfile)
    32  }
    33  
    34  // MaybeStartProfile checks the -blockProfile, -cpuProfile, and -memProfile flag and, for each that is set, attempts to start gathering profiling data into the appropriate files. It returns an object with one method, Stop(), that must be called in order to flush profile data to disk before the process terminates.
    35  func MaybeStartProfile() interface {
    36  	Stop()
    37  } {
    38  	p := &prof{}
    39  	if blockProfile != "" {
    40  		f, err := os.Create(blockProfile)
    41  		d.PanicIfError(err)
    42  		runtime.SetBlockProfileRate(1)
    43  		p.bp = f
    44  	}
    45  	if cpuProfile != "" {
    46  		f, err := os.Create(cpuProfile)
    47  		d.PanicIfError(err)
    48  		pprof.StartCPUProfile(f)
    49  		p.cpu = f
    50  	}
    51  	if memProfile != "" {
    52  		f, err := os.Create(memProfile)
    53  		d.PanicIfError(err)
    54  		p.mem = f
    55  	}
    56  	return p
    57  }
    58  
    59  type prof struct {
    60  	bp  io.WriteCloser
    61  	cpu io.Closer
    62  	mem io.WriteCloser
    63  }
    64  
    65  func (p *prof) Stop() {
    66  	if p.bp != nil {
    67  		pprof.Lookup("block").WriteTo(p.bp, 0)
    68  		p.bp.Close()
    69  		runtime.SetBlockProfileRate(0)
    70  	}
    71  	if p.cpu != nil {
    72  		pprof.StopCPUProfile()
    73  		p.cpu.Close()
    74  	}
    75  	if p.mem != nil {
    76  		pprof.WriteHeapProfile(p.mem)
    77  		p.mem.Close()
    78  	}
    79  }