1 files changed, 146 insertions, 0 deletions
diff --git a/src/runtime/pprof/mprof_test.go b/src/runtime/pprof/mprof_test.go
new file mode 100644
index 000000000..1880b9a70
--- /dev/null
+++ b/src/runtime/pprof/mprof_test.go
@@ -0,0 +1,146 @@
+// Copyright 2014 The Go Authors.  All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package pprof_test
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"regexp"
+	"runtime"
+	. "runtime/pprof"
+	"testing"
+	"unsafe"
+)
+
+var memSink interface{}
+
+func allocateTransient1M() {
+	for i := 0; i < 1024; i++ {
+		memSink = &struct{ x [1024]byte }{}
+	}
+}
+
+func allocateTransient2M() {
+	// prevent inlining
+	if memSink == nil {
+		panic("bad")
+	}
+	memSink = make([]byte, 2<<20)
+}
+
+type Obj32 struct {
+	link *Obj32
+	pad  [32 - unsafe.Sizeof(uintptr(0))]byte
+}
+
+var persistentMemSink *Obj32
+
+func allocatePersistent1K() {
+	for i := 0; i < 32; i++ {
+		// Can't use slice because that will introduce implicit allocations.
+		obj := &Obj32{link: persistentMemSink}
+		persistentMemSink = obj
+	}
+}
+
+var memoryProfilerRun = 0
+
+func TestMemoryProfiler(t *testing.T) {
+	// Create temp file for the profile.
+	f, err := ioutil.TempFile("", "memprof")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
+	defer func() {
+		f.Close()
+		os.Remove(f.Name())
+	}()
+
+	// Disable sampling, otherwise it's difficult to assert anything.
+	oldRate := runtime.MemProfileRate
+	runtime.MemProfileRate = 1
+	defer func() {
+		runtime.MemProfileRate = oldRate
+	}()
+	// Allocate a meg to ensure that mcache.next_sample is updated to 1.
+	for i := 0; i < 1024; i++ {
+		memSink = make([]byte, 1024)
+	}
+
+	// Do the interesting allocations.
+	allocateTransient1M()
+	allocateTransient2M()
+	allocatePersistent1K()
+	memSink = nil
+
+	runtime.GC() // materialize stats
+	if err := WriteHeapProfile(f); err != nil {
+		t.Fatalf("failed to write heap profile: %v", err)
+	}
+	f.Close()
+
+	memoryProfilerRun++
+	checkMemProfile(t, f.Name(), []string{"--alloc_space", "--show_bytes", "--lines"}, []string{
+		fmt.Sprintf(`%v .* runtime/pprof_test\.allocateTransient1M .*mprof_test.go:25`, 1<<20*memoryProfilerRun),
+		fmt.Sprintf(`%v .* runtime/pprof_test\.allocateTransient2M .*mprof_test.go:34`, 2<<20*memoryProfilerRun),
+		fmt.Sprintf(`%v .* runtime/pprof_test\.allocatePersistent1K .*mprof_test.go:47`, 1<<10*memoryProfilerRun),
+	}, []string{})
+
+	checkMemProfile(t, f.Name(), []string{"--inuse_space", "--show_bytes", "--lines"}, []string{
+		fmt.Sprintf(`%v .* runtime/pprof_test\.allocatePersistent1K .*mprof_test.go:47`, 1<<10*memoryProfilerRun),
+	}, []string{
+		"allocateTransient1M",
+		"allocateTransient2M",
+	})
+}
+
+func checkMemProfile(t *testing.T, file string, addArgs []string, what []string, whatnot []string) {
+	args := []string{"tool", "pprof", "--text"}
+	args = append(args, addArgs...)
+	args = append(args, os.Args[0], file)
+	out, err := exec.Command("go", args...).CombinedOutput()
+	if err != nil {
+		t.Fatalf("failed to execute pprof: %v\n%v\n", err, string(out))
+	}
+
+	matched := make(map[*regexp.Regexp]bool)
+	for _, s := range what {
+		matched[regexp.MustCompile(s)] = false
+	}
+	var not []*regexp.Regexp
+	for _, s := range whatnot {
+		not = append(not, regexp.MustCompile(s))
+	}
+
+	s := bufio.NewScanner(bytes.NewReader(out))
+	for s.Scan() {
+		ln := s.Text()
+		for re := range matched {
+			if re.MatchString(ln) {
+				if matched[re] {
+					t.Errorf("entry '%s' is matched twice", re.String())
+				}
+				matched[re] = true
+			}
+		}
+		for _, re := range not {
+			if re.MatchString(ln) {
+				t.Errorf("entry '%s' is matched, but must not", re.String())
+			}
+		}
+	}
+	for re, ok := range matched {
+		if !ok {
+			t.Errorf("entry '%s' is not matched", re.String())
+		}
+	}
+	if t.Failed() {
+		t.Logf("profile:\n%v", string(out))
+	}
+}