github.com/ssube/gitlab-ci-multi-runner@v1.2.1-0.20160607142738-b8d1285632e6/helpers/cli/cpuprofile.go (about)

     1  package cli_helpers
     2  
     3  import (
     4  	"os"
     5  	"runtime/pprof"
     6  
     7  	"github.com/codegangsta/cli"
     8  )
     9  
    10  func SetupCPUProfile(app *cli.App) {
    11  	app.Flags = append(app.Flags, cli.StringFlag{
    12  		Name:   "cpuprofile",
    13  		Usage:  "write cpu profile to file",
    14  		EnvVar: "CPU_PROFILE",
    15  	})
    16  
    17  	appBefore := app.Before
    18  	appAfter := app.After
    19  
    20  	app.Before = func(c *cli.Context) error {
    21  		if cpuProfile := c.String("cpuprofile"); cpuProfile != "" {
    22  			f, err := os.Create(cpuProfile)
    23  			if err != nil {
    24  				return err
    25  			}
    26  			pprof.StartCPUProfile(f)
    27  		}
    28  
    29  		if appBefore != nil {
    30  			return appBefore(c)
    31  		}
    32  		return nil
    33  	}
    34  
    35  	app.After = func(c *cli.Context) error {
    36  		pprof.StopCPUProfile()
    37  
    38  		if appAfter != nil {
    39  			return appAfter(c)
    40  		}
    41  		return nil
    42  	}
    43  }