Add MemProfileType to allow overriding type of memory profile

* Defaults to `heap`
* Allows setting to `allocs` to output allocation tracing profile

profile.go

@@ -43,6 +43,10 @@ type Profile struct {
 	// memProfileRate holds the rate for the memory profile.
 	memProfileRate int
 
+	// memProfileType holds the profile type for memory
+	// profiles. Allowed values are `heap` and `allocs`.
+	memProfileType string
+
 	// closer holds a cleanup function that run after each profile
 	closer func()
 
@@ -84,6 +88,19 @@ func MemProfileRate(rate int) func(*Profile) {
 	}
 }
 
+// MemProfileType changes which type of memory profiling to
+// enable. The only two memory types supported are alloc and heap. It
+// disables any previous profiling settings.
+func MemProfileType(memProfileType string) func(*Profile) {
+	if memProfileType != "heap" && memProfileType != "allocs" {
+		panic("Unexpected memory profile type, expected heap or alloc")
+	}
+	return func(p *Profile) {
+		p.memProfileType = memProfileType
+		p.mode = memMode
+	}
+}
+
 // MutexProfile enables mutex profiling.
 // It disables any previous profiling settings.
 func MutexProfile(p *Profile) { p.mode = mutexMode }
@@ -158,6 +175,10 @@ func Start(options ...func(*Profile)) interface {
 		}
 	}
 
+	if prof.memProfileType == "" {
+		prof.memProfileType = "heap"
+	}
+
 	switch prof.mode {
 	case cpuMode:
 		fn := filepath.Join(path, "cpu.pprof")
@@ -183,7 +204,7 @@ func Start(options ...func(*Profile)) interface {
 		runtime.MemProfileRate = prof.memProfileRate
 		logf("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, fn)
 		prof.closer = func() {
-			pprof.Lookup("heap").WriteTo(f, 0)
+			pprof.Lookup(prof.memProfileType).WriteTo(f, 0)
 			f.Close()
 			runtime.MemProfileRate = old
 			logf("profile: memory profiling disabled, %s", fn)