diff --git a/helpers/cli/cpuprofile.go b/helpers/cli/cpuprofile.go new file mode 100644 index 0000000000000000000000000000000000000000..45031852be5cfcbd035b73b989d9f63d2d34b268 --- /dev/null +++ b/helpers/cli/cpuprofile.go @@ -0,0 +1,45 @@ +package cli_helpers + +import ( + "os" + "runtime/pprof" + + "github.com/codegangsta/cli" +) + +func SetupCpuProfile(app *cli.App) { + app.Flags = append(app.Flags, cli.StringFlag{ + Name: "cpuprofile", + Usage: "write cpu profile to file", + EnvVar: "CPU_PROFILE", + }) + + appBefore := app.Before + appAfter := app.After + + app.Before = func(c *cli.Context) error { + if cpuProfile := c.String("cpuprofile"); cpuProfile != "" { + f, err := os.Create(cpuProfile) + if err != nil { + return err + } + pprof.StartCPUProfile(f) + } + + if appBefore != nil { + return appBefore(c) + } else { + return nil + } + } + + app.After = func(c *cli.Context) error { + pprof.StopCPUProfile() + + if appAfter != nil { + return appAfter(c) + } else { + return nil + } + } +} diff --git a/main.go b/main.go index e581a98309055804a8da7a7d1dff57d1823d817e..acba16573458183c5e848fc7c28e86c471d44c22 100644 --- a/main.go +++ b/main.go @@ -60,6 +60,7 @@ func main() { }, } cli_helpers.SetupLogLevelOptions(app) + cli_helpers.SetupCpuProfile(app) app.Commands = common.GetCommands() app.CommandNotFound = func(context *cli.Context, command string) { logrus.Fatalln("Command", command, "not found.")