tools/benchmark: add flags for pprof to storage put

This commit adds flags for profiling with runtime/pprof to storage
put:
- --cpuprof-path: specify a path of CPU profiling result, if it is not
    empty, profiling is activated
- --memprof-path: specify a path of heap profiling result, if it is not
    empty, profiling is activated

Of course, the flags should be added to RootCmd ideally. However,
adding common flags that shared by children command requires the
ongoing PR: spf13/cobra#220 . Therefore this
commit adds the flags to storage put only.

tools/benchmark/cmd/root.go

@@ -41,6 +41,9 @@ var (
 	bar     *pb.ProgressBar
 	results chan result
 	wg      sync.WaitGroup
+
+	cpuProfPath string
+	memProfPath string
 )
 
 func init() {

tools/benchmark/cmd/storage-put.go

@@ -18,6 +18,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"os"
+	"runtime/pprof"
 	"time"
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
@@ -46,6 +47,11 @@ func init() {
 	storagePutCmd.Flags().IntVar(&storageKeySize, "key-size", 64, "a size of key (Byte)")
 	storagePutCmd.Flags().IntVar(&valueSize, "value-size", 64, "a size of value (Byte)")
 	storagePutCmd.Flags().BoolVar(&txn, "txn", false, "put a key in transaction or not")
+
+	// TODO: after the PR https://github.com/spf13/cobra/pull/220 is merged, the below pprof related flags should be moved to RootCmd
+	storagePutCmd.Flags().StringVar(&cpuProfPath, "cpuprof-path", "", "a path of file for storing cpu profile result")
+	storagePutCmd.Flags().StringVar(&memProfPath, "memprof-path", "", "a path of file for storing heap profile result")
+
 }
 
 func createBytesSlice(bytesN, sliceN int) [][]byte {
@@ -60,6 +66,37 @@ func createBytesSlice(bytesN, sliceN int) [][]byte {
 }
 
 func storagePutFunc(cmd *cobra.Command, args []string) {
+	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
+			}
+		}()
+	}
+
 	keys := createBytesSlice(storageKeySize, totalNrKeys)
 	vals := createBytesSlice(valueSize, totalNrKeys)