Merge pull request #54 from AeroNotix/allow-override-memory-profile-type

Add MemProfileType to allow overriding type of memory profile

Author: davecheney

README.md

@@ -28,7 +28,7 @@ func main() {
 options
 -------
 
-What to profile is controlled by config value passed to profile.Start. 
+What to profile is controlled by config value passed to profile.Start.
 By default CPU profiling is enabled.
 
 ```go
@@ -39,6 +39,9 @@ func main() {
     // ensure profiling information is written to disk.
     p := profile.Start(profile.MemProfile, profile.ProfilePath("."), profile.NoShutdownHook)
     ...
+    // You can enable different kinds of memory profiling, either Heap or Allocs where Heap
+    // profiling is the default with profile.MemProfile.
+    p := profile.Start(profile.MemProfileAllocs, profile.ProfilePath("."), profile.NoShutdownHook)
 }
 ```
 

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,24 @@ func MemProfileRate(rate int) func(*Profile) {
 	}
 }
 
+// MemProfileHeap changes which type of memory profiling to profile
+// the heap.
+func MemProfileHeap() func(*Profile) {
+	return func(p *Profile) {
+		p.memProfileType = "heap"
+		p.mode = memMode
+	}
+}
+
+// MemProfileAllocs changes which type of memory to profile
+// allocations.
+func MemProfileAllocs() func(*Profile) {
+	return func(p *Profile) {
+		p.memProfileType = "allocs"
+		p.mode = memMode
+	}
+}
+
 // MutexProfile enables mutex profiling.
 // It disables any previous profiling settings.
 func MutexProfile(p *Profile) { p.mode = mutexMode }
@@ -158,6 +180,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 +209,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)