From 1c802e9266f4d203fd6676338338d8e8e159d9b7 Mon Sep 17 00:00:00 2001
From: Hitoshi Mitake <mitake.hitoshi@lab.ntt.co.jp>
Date: Fri, 8 Jan 2016 17:31:28 +0900
Subject: [PATCH] tools/benchmark: add flags for pprof to storage put

This commit adds flags for profiling with runtime/pprof to storage
put:
- --cpuprofile: specify a path of CPU profiling result, if it is not
    empty, profiling is activated
- --memprofile: 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: https://github.com/spf13/cobra/pull/220 . Therefore this
commit adds the flags to storage put only.
---
 tools/benchmark/cmd/root.go        |  3 +++
 tools/benchmark/cmd/storage-put.go | 37 ++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/tools/benchmark/cmd/root.go b/tools/benchmark/cmd/root.go
index 8283ea3a3..a7a00fecd 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 e4e41e75a..f715b0fa9 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, "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)