From 5df85770fca6184e4230e9cb4fddeb841087563f Mon Sep 17 00:00:00 2001
From: bep <bjorn.erik.pedersen@gmail.com>
Date: Thu, 5 Feb 2015 15:48:09 +0100
Subject: [PATCH] Add memprofile to pprof benchmark

---
 commands/benchmark.go | 41 +++++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/commands/benchmark.go b/commands/benchmark.go
index d67ec9cba..1e0cc8d7b 100644
--- a/commands/benchmark.go
+++ b/commands/benchmark.go
@@ -20,6 +20,7 @@ import (
 )
 
 var cpuProfilefile string
+var memProfilefile string
 var benchmarkTimes int
 
 var benchmark = &cobra.Command{
@@ -34,21 +35,41 @@ running process creating a benchmark.`,
 }
 
 func init() {
-	benchmark.Flags().StringVar(&cpuProfilefile, "outputfile", "/tmp/hugo-cpuprofile", "path/filename for the profile file")
+	benchmark.Flags().StringVar(&cpuProfilefile, "cpuprofile", "", "path/filename for the CPU profile file")
+	benchmark.Flags().StringVar(&memProfilefile, "memprofile", "", "path/filename for the memory profile file")
+
 	benchmark.Flags().IntVarP(&benchmarkTimes, "count", "n", 13, "number of times to build the site")
 }
 
 func bench(cmd *cobra.Command, args []string) {
-	f, err := os.Create(cpuProfilefile)
 
-	if err != nil {
-		panic(err)
+	if memProfilefile != "" {
+		f, err := os.Create(memProfilefile)
+
+		if err != nil {
+			panic(err)
+		}
+		for i := 0; i < benchmarkTimes; i++ {
+			_ = buildSite()
+		}
+		pprof.WriteHeapProfile(f)
+		f.Close()
+
+	} else {
+		if cpuProfilefile == "" {
+			cpuProfilefile = "/tmp/hugo-cpuprofile"
+		}
+		f, err := os.Create(cpuProfilefile)
+
+		if err != nil {
+			panic(err)
+		}
+
+		pprof.StartCPUProfile(f)
+		defer pprof.StopCPUProfile()
+		for i := 0; i < benchmarkTimes; i++ {
+			_ = buildSite()
+		}
 	}
 
-	pprof.StartCPUProfile(f)
-	defer pprof.StopCPUProfile()
-
-	for i := 0; i < benchmarkTimes; i++ {
-		_ = buildSite()
-	}
 }