testing: add -test.memprofile and -test.memprofilerate flags.

These allow a test to generate memory profiles automatically.

R=golang-dev, rsc1
CC=golang-dev
https://golang.org/cl/4273064

src/cmd/gotest/doc.go

@@ -38,11 +38,11 @@ interfere with the non-test installation.
 Usage:
 	gotest [pkg_test.go ...]
 
-The resulting binary, called (for amd64) 6.out, has a couple of
-arguments.
+The resulting binary, called (for amd64) 6.out, has several flags.
 
 Usage:
-	6.out [-test.v] [-test.run pattern] [-test.bench pattern]
+	6.out [-test.v] [-test.run pattern] [-test.bench pattern] \
+		[test.memprofile=prof.out] [-test.memprofilerate=1]
 
 The -test.v flag causes the tests to be logged as they run.  The
 -test.run flag causes only those tests whose names match the regular
@@ -52,5 +52,19 @@ exit code.  If any tests fail, it prints FAIL and exits with a
 non-zero code.  The -test.bench flag is analogous to the -test.run
 flag, but applies to benchmarks.  No benchmarks run by default.
 
+The -test.memprofile flag causes the testing software to write a
+memory profile to the specified file when all tests are complete.  Use
+-test.run or -test.bench to limit the profile to a particular test or
+benchmark.  The -test.memprofilerate flag enables more precise (and
+expensive) profiles by setting runtime.MemProfileRate;
+	godoc runtime MemProfileRate
+for details.  The defaults are no memory profile and the standard
+setting of MemProfileRate.  The memory profile records a sampling of
+the memory in use at the end of the test.  To profile all memory
+allocations, use -test.memprofilerate=1 to sample every byte and set
+the environment variable GOGC=off to disable the garbage collector,
+provided the test can run in the available memory without garbage
+collection.
+
 */
 package documentation

src/pkg/gob/timing_test.go

@@ -6,12 +6,10 @@ package gob
 
 import (
 	"bytes"
-	"flag"
 	"fmt"
 	"io"
 	"os"
 	"runtime"
-	"runtime/pprof"
 	"testing"
 )
 
@@ -22,8 +20,6 @@ type Bench struct {
 	D []byte
 }
 
-var memprofile = flag.String("memprofile", "", "write the memory profile in Test*Mallocs to the named file")
-
 func benchmarkEndToEnd(r io.Reader, w io.Writer, b *testing.B) {
 	b.StopTimer()
 	enc := NewEncoder(w)
@@ -54,7 +50,6 @@ func BenchmarkEndToEndByteBuffer(b *testing.B) {
 }
 
 func TestCountEncodeMallocs(t *testing.T) {
-	runtime.MemProfileRate = 1
 	var buf bytes.Buffer
 	enc := NewEncoder(&buf)
 	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}
@@ -67,21 +62,10 @@ func TestCountEncodeMallocs(t *testing.T) {
 		}
 	}
 	mallocs += runtime.MemStats.Mallocs
-	if *memprofile != "" {
-		if fd, err := os.Open(*memprofile, os.O_WRONLY|os.O_CREAT|os.O_TRUNC, 0666); err != nil {
-			t.Errorf("can't open %s: %s", *memprofile, err)
-		} else {
-			if err = pprof.WriteHeapProfile(fd); err != nil {
-				t.Errorf("can't write %s: %s", *memprofile, err)
-			}
-			fd.Close()
-		}
-	}
 	fmt.Printf("mallocs per encode of type Bench: %d\n", mallocs/count)
 }
 
 func TestCountDecodeMallocs(t *testing.T) {
-	runtime.MemProfileRate = 1
 	var buf bytes.Buffer
 	enc := NewEncoder(&buf)
 	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}
@@ -102,15 +86,5 @@ func TestCountDecodeMallocs(t *testing.T) {
 		}
 	}
 	mallocs += runtime.MemStats.Mallocs
-	if *memprofile != "" {
-		if fd, err := os.Open(*memprofile, os.O_WRONLY|os.O_CREAT|os.O_TRUNC, 0666); err != nil {
-			t.Errorf("can't open %s: %s", *memprofile, err)
-		} else {
-			if err = pprof.WriteHeapProfile(fd); err != nil {
-				t.Errorf("can't write %s: %s", *memprofile, err)
-			}
-			fd.Close()
-		}
-	}
 	fmt.Printf("mallocs per decode of type Bench: %d\n", mallocs/count)
 }

src/pkg/rpc/server_test.go

@@ -5,14 +5,12 @@
 package rpc
 
 import (
-	"flag"
 	"fmt"
 	"http/httptest"
 	"log"
 	"net"
 	"os"
 	"runtime"
-	"runtime/pprof"
 	"strings"
 	"sync"
 	"testing"
@@ -25,8 +23,6 @@ var (
 	once, newOnce, httpOnce   sync.Once
 )
 
-var memprofile = flag.String("memprofile", "", "write the memory profile in TestCountMallocs to the named file")
-
 const (
 	second      = 1e9
 	newHttpPath = "/foo"
@@ -356,7 +352,6 @@ func testSendDeadlock(client *Client) {
 }
 
 func TestCountMallocs(t *testing.T) {
-	runtime.MemProfileRate = 1
 	once.Do(startServer)
 	client, err := Dial("tcp", serverAddr)
 	if err != nil {
@@ -365,7 +360,7 @@ func TestCountMallocs(t *testing.T) {
 	args := &Args{7, 8}
 	reply := new(Reply)
 	mallocs := 0 - runtime.MemStats.Mallocs
-	const count = 10000
+	const count = 100
 	for i := 0; i < count; i++ {
 		err = client.Call("Arith.Add", args, reply)
 		if err != nil {
@@ -376,16 +371,6 @@ func TestCountMallocs(t *testing.T) {
 		}
 	}
 	mallocs += runtime.MemStats.Mallocs
-	if *memprofile != "" {
-		if fd, err := os.Open(*memprofile, os.O_WRONLY|os.O_CREAT|os.O_TRUNC, 0666); err != nil {
-			t.Errorf("can't open %s: %s", *memprofile, err)
-		} else {
-			if err = pprof.WriteHeapProfile(fd); err != nil {
-				t.Errorf("can't write %s: %s", *memprofile, err)
-			}
-			fd.Close()
-		}
-	}
 	fmt.Printf("mallocs per rpc round trip: %d\n", mallocs/count)
 }
 

src/pkg/testing/testing.go

@@ -43,12 +43,17 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"runtime/pprof"
 	"time"
 )
 
-// Report as tests are run; default is silent for success.
-var chatty = flag.Bool("test.v", false, "verbose: print additional output")
-var match = flag.String("test.run", "", "regular expression to select tests to run")
+var (
+	// Report as tests are run; default is silent for success.
+	chatty         = flag.Bool("test.v", false, "verbose: print additional output")
+	match          = flag.String("test.run", "", "regular expression to select tests to run")
+	memProfile     = flag.String("test.memprofile", "", "after execution write the memory profile to the named file")
+	memProfileRate = flag.Int("test.memprofilerate", 0, "if >=0, sets runtime.MemProfileRate")
+)
 
 
 // Insert final newline if needed and tabs after internal newlines.
@@ -138,6 +143,8 @@ func tRunner(t *T, test *InternalTest) {
 // of gotest.
 func Main(matchString func(pat, str string) (bool, os.Error), tests []InternalTest) {
 	flag.Parse()
+
+	before()
 	ok := true
 	if len(tests) == 0 {
 		println("testing: warning: no tests to run")
@@ -170,9 +177,33 @@ func Main(matchString func(pat, str string) (bool, os.Error), tests []InternalTe
 			print(t.errors)
 		}
 	}
+	after()
 	if !ok {
 		println("FAIL")
 		os.Exit(1)
 	}
 	println("PASS")
 }
+
+// before runs before all testing.
+func before() {
+	if *memProfileRate > 0 {
+		runtime.MemProfileRate = *memProfileRate
+	}
+}
+
+// after runs after all testing.
+func after() {
+	if *memProfile == "" {
+		return
+	}
+	fd, err := os.Open(*memProfile, os.O_WRONLY|os.O_CREAT|os.O_TRUNC, 0666)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "testing: can't open %s: %s", *memProfile, err)
+		return
+	}
+	if err = pprof.WriteHeapProfile(fd); err != nil {
+		fmt.Fprintf(os.Stderr, "testing: can't write %s: %s", *memProfile, err)
+	}
+	fd.Close()
+}