Merge pull request #4201 from mitake/benchmark-pprof

tools/benchmark: add flags for pprof to storage put
This commit is contained in:
Xiang Li 2016-01-13 20:17:30 -08:00
commit 6c82d768b2
2 changed files with 40 additions and 0 deletions

View File

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

View File

@ -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)