Merge pull request #4201 from mitake/benchmark-pprof

tools/benchmark: add flags for pprof to storage put

tools/benchmark/cmd/root.go

@@ -39,6 +39,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, "cpuprofile", "", "the path of file for storing cpu profile result")
+	storagePutCmd.Flags().StringVar(&memProfPath, "memprofile", "", "the 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)