cmd/evm: added mem/cpu profiling

2 files changed, 44 insertions, 0 deletions

diff --git a/cmd/evm/main.go b/cmd/evm/main.go
index 5f85f484e..e85d31d03 100644
--- a/cmd/evm/main.go
+++ b/cmd/evm/main.go
@@ -35,6 +35,18 @@ var (
 		Name:  "debug",
 		Usage: "output full trace logs",
 	}
+	MemProfileFlag = cli.StringFlag{
+		Name:  "memprofile",
+		Usage: "creates a memory profile at the given path",
+	}
+	CPUProfileFlag = cli.StringFlag{
+		Name:  "cpuprofile",
+		Usage: "creates a CPU profile at the given path",
+	}
+	StatDumpFlag = cli.BoolFlag{
+		Name:  "statdump",
+		Usage: "displays stack and heap memory information",
+	}
 	CodeFlag = cli.StringFlag{
 		Name:  "code",
 		Usage: "EVM code",
@@ -93,6 +105,9 @@ func init() {
 		DumpFlag,
 		InputFlag,
 		DisableGasMeteringFlag,
+		MemProfileFlag,
+		CPUProfileFlag,
+		StatDumpFlag,
 	}
 	app.Commands = []cli.Command{
 		compileCommand,
diff --git a/cmd/evm/runner.go b/cmd/evm/runner.go
index 6ef9230f4..22538d7b1 100644
--- a/cmd/evm/runner.go
+++ b/cmd/evm/runner.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"runtime/pprof"
 	"time"
 
 	goruntime "runtime"
@@ -108,6 +109,19 @@ func runCmd(ctx *cli.Context) error {
 		},
 	}
 
+	if cpuProfilePath := ctx.GlobalString(CPUProfileFlag.Name); cpuProfilePath != "" {
+		f, err := os.Create(cpuProfilePath)
+		if err != nil {
+			fmt.Println("could not create CPU profile: ", err)
+			os.Exit(1)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			fmt.Println("could not start CPU profile: ", err)
+			os.Exit(1)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
 	tstart := time.Now()
 	if ctx.GlobalBool(CreateFlag.Name) {
 		input := append(code, common.Hex2Bytes(ctx.GlobalString(InputFlag.Name))...)
@@ -125,12 +139,27 @@ func runCmd(ctx *cli.Context) error {
 		fmt.Println(string(statedb.Dump()))
 	}
 
+	if memProfilePath := ctx.GlobalString(MemProfileFlag.Name); memProfilePath != "" {
+		f, err := os.Create(memProfilePath)
+		if err != nil {
+			fmt.Println("could not create memory profile: ", err)
+			os.Exit(1)
+		}
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			fmt.Println("could not write memory profile: ", err)
+			os.Exit(1)
+		}
+		f.Close()
+	}
+
 	if ctx.GlobalBool(DebugFlag.Name) {
 		fmt.Fprintln(os.Stderr, "#### TRACE ####")
 		vm.WriteTrace(os.Stderr, logger.StructLogs())
 		fmt.Fprintln(os.Stderr, "#### LOGS ####")
 		vm.WriteLogs(os.Stderr, statedb.Logs())
+	}
 
+	if ctx.GlobalBool(StatDumpFlag.Name) {
 		var mem goruntime.MemStats
 		goruntime.ReadMemStats(&mem)
 		fmt.Fprintf(os.Stderr, `evm execution time: %v