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