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