Add -cpuprofile and -memprofile to dexcon-simulation (#56)

1 files changed, 28 insertions, 0 deletions

diff --git a/cmd/dexcon-simulation/main.go b/cmd/dexcon-simulation/main.go
index 5547700..def4ac9 100644
--- a/cmd/dexcon-simulation/main.go
+++ b/cmd/dexcon-simulation/main.go
@@ -20,8 +20,11 @@ package main
 import (
 	"flag"
 	"fmt"
+	"log"
 	"math/rand"
 	"os"
+	"runtime"
+	"runtime/pprof"
 	"time"
 
 	"github.com/dexon-foundation/dexon-consensus-core/simulation"
@@ -30,6 +33,8 @@ import (
 
 var initialize = flag.Bool("init", false, "initialize config file")
 var configFile = flag.String("config", "", "path to simulation config file")
+var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")
+var memprofile = flag.String("memprofile", "", "write memory profile to `file`")
 
 func main() {
 	flag.Parse()
@@ -49,5 +54,28 @@ func main() {
 		//os.Exit(0)
 	}
 
+	if *cpuprofile != "" {
+		f, err := os.Create(*cpuprofile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
 	simulation.Run(*configFile)
+
+	if *memprofile != "" {
+		f, err := os.Create(*memprofile)
+		if err != nil {
+			log.Fatal("could not create memory profile: ", err)
+		}
+		runtime.GC() // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("could not write memory profile: ", err)
+		}
+		f.Close()
+	}
 }