cmd/starlark: add Go memory profiling support (#190)

* cmd/starlark: add Go memory profiling support

For #184

* cmd/starlark: check errors in profiling code

1 files changed, 32 insertions, 16 deletions

diff --git a/cmd/starlark/starlark.go b/cmd/starlark/starlark.go
index 2ae6c71..58774dc 100644
--- a/cmd/starlark/starlark.go
+++ b/cmd/starlark/starlark.go
@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"runtime"
 	"runtime/pprof"
 	"strings"
 
@@ -22,6 +23,7 @@ import (
 // flags
 var (
 	cpuprofile = flag.String("cpuprofile", "", "gather Go CPU profile in this file")
+	memprofile = flag.String("memprofile", "", "gather Go memory profile in this file")
 	profile    = flag.String("profile", "", "gather Starlark time profile in this file")
 	showenv    = flag.Bool("showenv", false, "on success, print final global environment")
 	execprog   = flag.String("c", "", "execute program `prog`")
@@ -48,27 +50,35 @@ func doMain() int {
 
 	if *cpuprofile != "" {
 		f, err := os.Create(*cpuprofile)
-		if err != nil {
-			log.Fatal(err)
-		}
-		if err := pprof.StartCPUProfile(f); err != nil {
-			log.Fatal(err)
-		}
-		defer pprof.StopCPUProfile()
+		check(err)
+		err = pprof.StartCPUProfile(f)
+		check(err)
+		defer func() {
+			pprof.StopCPUProfile()
+			err := f.Close()
+			check(err)
+		}()
+	}
+	if *memprofile != "" {
+		f, err := os.Create(*memprofile)
+		check(err)
+		defer func() {
+			runtime.GC()
+			err := pprof.Lookup("heap").WriteTo(f, 0)
+			check(err)
+			err = f.Close()
+			check(err)
+		}()
 	}
 
 	if *profile != "" {
 		f, err := os.Create(*profile)
-		if err != nil {
-			log.Fatal(err)
-		}
-		if err := starlark.StartProfile(f); err != nil {
-			log.Fatal(err)
-		}
+		check(err)
+		err = starlark.StartProfile(f)
+		check(err)
 		defer func() {
-			if err := starlark.StopProfile(); err != nil {
-				log.Fatal(err)
-			}
+			err := starlark.StopProfile()
+			check(err)
 		}()
 	}
 
@@ -117,3 +127,9 @@ func doMain() int {
 
 	return 0
 }
+
+func check(err error) {
+	if err != nil {
+		log.Fatal(err)
+	}
+}