Merge pull request #37 from PatrickCronin/pcronin/mem-profile

Facilitate mem profiling in the feature branch

Author: PatrickCronin

cmd/routesum/main.go

@@ -15,7 +15,8 @@ import (
 )
 
 func main() {
-	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
+	cpuProfile := flag.String("cpuprofile", "", "write cpu profile to file")
+	memProfile := flag.String("memprofile", "", "write mem profile to file")
 
 	showMemStats := flag.Bool(
 		"show-mem-stats",
@@ -30,15 +31,23 @@ func main() {
 	}
 
 	var cpuProfOut io.Writer
-	if *cpuprofile != "" {
+	if *cpuProfile != "" {
 		var err error
-		if cpuProfOut, err = os.Create(*cpuprofile); err != nil {
+		if cpuProfOut, err = os.Create(*cpuProfile); err != nil {
 			fmt.Fprint(os.Stderr, errors.Wrap(err, "create cpu profile output file").Error())
 			os.Exit(1)
 		}
 	}
 
-	if err := summarize(os.Stdin, os.Stdout, memStatsOut, cpuProfOut); err != nil {
+	var memProfOut io.WriteCloser
+	if *memProfile != "" {
+		var err error
+		if memProfOut, err = os.Create(*memProfile); err != nil {
+			fmt.Fprint(os.Stderr, errors.Wrap(err, "create mem profile output file").Error())
+		}
+	}
+
+	if err := summarize(os.Stdin, os.Stdout, memStatsOut, cpuProfOut, memProfOut); err != nil {
 		fmt.Fprintf(os.Stderr, "summarize: %s\n", err.Error())
 		os.Exit(1)
 	}
@@ -47,6 +56,7 @@ func main() {
 func summarize(
 	in io.Reader,
 	out, memStatsOut, cpuProfOut io.Writer,
+	memProfOut io.WriteCloser,
 ) error {
 	if cpuProfOut != nil {
 		if err := pprof.StartCPUProfile(cpuProfOut); err != nil {
@@ -68,6 +78,15 @@ func summarize(
 		}
 	}
 
+	if memProfOut != nil {
+		if err := pprof.WriteHeapProfile(memProfOut); err != nil {
+			return errors.Wrap(err, "write mem profile")
+		}
+		if err := memProfOut.Close(); err != nil {
+			return errors.Wrap(err, "close mem profile")
+		}
+	}
+
 	if memStatsOut != nil {
 		numInternalNodes, numLeafNodes, internalNodesTotalSize, leafNodesTotalSize := rs.MemUsage()
 		fmt.Fprintf(memStatsOut,

cmd/routesum/main_test.go

@@ -41,7 +41,7 @@ func TestSummarize(t *testing.T) {
 				memStatsOut = &memStatsBuilder
 			}
 
-			err := summarize(in, &out, memStatsOut, nil)
+			err := summarize(in, &out, memStatsOut, nil, nil)
 			require.NoError(t, err, "summarize does not throw an error")
 
 			assert.Equal(t, "192.0.2.0/31\n", out.String(), "read expected output")