Add memprofile to pprof benchmark

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 {
+	if memProfilefile != "" {
-		panic(err)
+		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()
-	}
 }