Add CPU profiling flag

1 files changed, 23 insertions, 11 deletions

diff --git a/cmd/planr/main.go b/cmd/planr/main.go
index 4ff57eb..d7f62aa 100644
--- a/cmd/planr/main.go
+++ b/cmd/planr/main.go
@@ -1,15 +1,16 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"io"
 	"log"
 	"os"
-        "flag"
+	"runtime/pprof"
 
 	"golang.flu0r1ne.net/planr"
-	"golang.flu0r1ne.net/planr/cmd/planr/sub"
 	"golang.flu0r1ne.net/planr/adapters/gtest"
+	"golang.flu0r1ne.net/planr/cmd/planr/sub"
 )
 
 func printUsage(w io.Writer) {
@@ -27,6 +28,11 @@ func dieUsage() {
 	os.Exit(1)
 }
 
+var src        = flag.String("srcdir", "", "source directory")
+var config     = flag.String("configdir", "", "config directory")
+var build      = flag.String("builddir", "", "build directory")
+var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
+
 func getConfiguredRunner() planr.Runner {
   r := planr.ConfigureRunner()
   r  = planr.RegisterAdapter(r, gtest.NewAdapter())
@@ -35,21 +41,15 @@ func getConfiguredRunner() planr.Runner {
     r = planr.SetConfigDirFromTree(r, wd) 
   }
 
-  src    := flag.String("srcdir", "", "source directory")
-  config := flag.String("configdir", "", "config directory")
-  build  := flag.String("builddir", "", "build directory")
-
-  flag.Parse()
-
-  if src != nil && *src != "" {
+  if *src != "" {
     r = planr.SetSrcDir(r, *src)    
   }
 
-  if config != nil && *config != "" {
+  if *config != "" {
     r = planr.SetConfigDir(r, *config)
   }
 
-  if build != nil && *build != "" {
+  if *build != "" {
     r = planr.SetBuildDir(r, *build)
   }
 
@@ -57,6 +57,18 @@ func getConfiguredRunner() planr.Runner {
 }
 
 func main() {
+        flag.Parse()
+
+        if *cpuprofile != "" {
+          f, err := os.Create(*cpuprofile)
+          
+          if err != nil {
+            log.Fatal(err)
+          }
+
+          pprof.StartCPUProfile(f)
+          defer pprof.StopCPUProfile()
+        }
 
         log.SetFlags(log.Llongfile | log.Lmsgprefix)
         log.SetPrefix("planr: ")