1 files changed, 44 insertions, 0 deletions
diff --git a/src/cmd/internal/gc/util.go b/src/cmd/internal/gc/util.go
index ceb3eead35..df68d50058 100644
--- a/src/cmd/internal/gc/util.go
+++ b/src/cmd/internal/gc/util.go
@@ -2,6 +2,8 @@ package gc
 
 import (
 	"cmd/internal/obj"
+	"os"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 )
@@ -68,3 +70,45 @@ func stringsCompare(a, b string) int {
 	}
 	return +1
 }
+
+var atExitFuncs []func()
+
+func AtExit(f func()) {
+	atExitFuncs = append(atExitFuncs, f)
+}
+
+func Exit(code int) {
+	for i := len(atExitFuncs) - 1; i >= 0; i-- {
+		f := atExitFuncs[i]
+		atExitFuncs = atExitFuncs[:i]
+		f()
+	}
+	os.Exit(code)
+}
+
+var cpuprofile string
+var memprofile string
+
+func startProfile() {
+	if cpuprofile != "" {
+		f, err := os.Create(cpuprofile)
+		if err != nil {
+			Fatal("%v", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			Fatal("%v", err)
+		}
+		AtExit(pprof.StopCPUProfile)
+	}
+	if memprofile != "" {
+		f, err := os.Create(memprofile)
+		if err != nil {
+			Fatal("%v", err)
+		}
+		AtExit(func() {
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				Fatal("%v", err)
+			}
+		})
+	}
+}