Add memprofile to pprof benchmark

diff --git a/commands/benchmark.go b/commands/benchmark.go
index d67ec9cba110af9fe3452d1f54643f5a6989c07b..1e0cc8d7bae29ee7ca56f60c9fa12a6403fecc96 100644 (file)
--- 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()
 
-       pprof.StartCPUProfile(f)
-       defer pprof.StopCPUProfile()
+       } else {
+               if cpuProfilefile == "" {
+                       cpuProfilefile = "/tmp/hugo-cpuprofile"
+               }
+               f, err := os.Create(cpuProfilefile)
 
-       for i := 0; i < benchmarkTimes; i++ {
-               _ = buildSite()
+               if err != nil {
+                       panic(err)
+               }
+
+               pprof.StartCPUProfile(f)
+               defer pprof.StopCPUProfile()
+               for i := 0; i < benchmarkTimes; i++ {
+                       _ = buildSite()
+               }
        }
+
 }