Added --cpuprofile flag for the dev Docker images

cmd/root_profiling.go

@@ -0,0 +1,55 @@
+// +build profiling
+
+package cmd
+
+import (
+	"log"
+	"os"
+	"runtime/pprof"
+
+	"github.com/spf13/cobra"
+)
+
+func init() {
+	var profilePath string
+	RootCmd.PersistentFlags().StringVar(&profilePath, "cpuprofile", "", "enables pprof profiling and sets its output path")
+
+	pprofEnabled := false
+	originalPersistentPreRunE := RootCmd.PersistentPreRunE
+	RootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		if profilePath == "" {
+			return nil
+		}
+
+		f, err := os.Create(profilePath)
+		if err != nil {
+			return err
+		}
+
+		log.Println("enabling profiling")
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			return err
+		}
+
+		pprofEnabled = true
+
+		if originalPersistentPreRunE != nil {
+			return originalPersistentPreRunE(cmd, args)
+		}
+
+		return nil
+	}
+
+	originalPersistentPostRun := RootCmd.PersistentPreRun
+	RootCmd.PersistentPostRun = func(cmd *cobra.Command, args []string) {
+		if pprofEnabled {
+			log.Println("shutting down profiling")
+			pprof.StopCPUProfile()
+		}
+
+		if originalPersistentPostRun != nil {
+			originalPersistentPostRun(cmd, args)
+		}
+	}
+}