diff --git a/tools/benchmark/cmd/root.go b/tools/benchmark/cmd/root.go index 8283ea3a391d..a7a00fecd602 100644 --- a/tools/benchmark/cmd/root.go +++ b/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() { diff --git a/tools/benchmark/cmd/storage-put.go b/tools/benchmark/cmd/storage-put.go index e4e41e75a6e2..4eabaf9f121f 100644 --- a/tools/benchmark/cmd/storage-put.go +++ b/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)