Skip to content

Commit

Permalink
Merge pull request #4201 from mitake/benchmark-pprof
Browse files Browse the repository at this point in the history
tools/benchmark: add flags for pprof to storage put
  • Loading branch information
xiang90 committed Jan 14, 2016
2 parents a8a7865 + 1c802e9 commit 6c82d76
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tools/benchmark/cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,9 @@ var (
bar *pb.ProgressBar
results chan result
wg sync.WaitGroup

cpuProfPath string
memProfPath string
)

func init() {
Expand Down
37 changes: 37 additions & 0 deletions tools/benchmark/cmd/storage-put.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"crypto/rand"
"fmt"
"os"
"runtime/pprof"
"time"

"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
Expand Down Expand Up @@ -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 {
Expand All @@ -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)

Expand Down

0 comments on commit 6c82d76

Please sign in to comment.