wip

ToDo: use this spf13/cobra#220

tools/benchmark/cmd/root.go

@@ -17,6 +17,7 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"runtime/pprof"
 	"sync"
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
@@ -41,15 +42,51 @@ var (
 	bar     *pb.ProgressBar
 	results chan result
 	wg      sync.WaitGroup
+
+	cpuProfPath string
+	memProfPath string
 )
 
 func init() {
 	RootCmd.PersistentFlags().StringVar(&endpoints, "endpoint", "127.0.0.1:2378", "comma-separated gRPC endpoints")
 	RootCmd.PersistentFlags().UintVar(&totalConns, "conns", 1, "Total number of gRPC connections")
 	RootCmd.PersistentFlags().UintVar(&totalClients, "clients", 1, "Total number of gRPC clients")
+	RootCmd.PersistentFlags().StringVar(&cpuProfPath, "cpuprof-path", "", "a path of file for storing cpu profile result")
+	RootCmd.PersistentFlags().StringVar(&memProfPath, "memprof-path", "", "a path of file for storing heap profile result")
 }
 
 func Execute() {
+	if cpuProfPath != "" {
+		f, err := os.Create(cpuProfPath)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "failed to create a file for storing cpu profile result: ", err)
+			os.Exit(1)
+		}
+
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "failed to start cpu profile: ", err)
+			os.Exit(1)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
+	if memProfPath != "" {
+		f, err := os.Create(memProfPath)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "failed to create a file for storing heap profile result: ", err)
+			os.Exit(1)
+		}
+
+		defer func() {
+			err := pprof.WriteHeapProfile(f)
+			if err != nil {
+				fmt.Fprintln(os.Stderr, "failed to write heap profile result: ", err)
+				// can do nothing for handling the error
+			}
+		}()
+	}
+
 	if err := RootCmd.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(-1)