From 411b9330bdfe74ef9b8b47464c9db57e6787b3b5 Mon Sep 17 00:00:00 2001 From: Kamil Trzcinski <ayufan@ayufan.eu> Date: Sat, 30 Jan 2016 20:11:43 +0100 Subject: [PATCH] Add CPU profile option --- helpers/cli/cpuprofile.go | 45 +++++++++++++++++++++++++++++++++++++++ main.go | 1 + 2 files changed, 46 insertions(+) create mode 100644 helpers/cli/cpuprofile.go diff --git a/helpers/cli/cpuprofile.go b/helpers/cli/cpuprofile.go new file mode 100644 index 000000000..45031852b --- /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 e581a9830..acba16573 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.") -- GitLab