testing: add memory allocation stats to benchmark

R=rsc, nigeltao, dave, bradfitz, r, rogpeppe
CC=golang-dev
https://golang.org/cl/6497084

src/pkg/testing/benchmark.go

@@ -9,11 +9,19 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"sync"
 	"time"
 )
 
 var matchBenchmarks = flag.String("test.bench", "", "regular expression to select benchmarks to run")
 var benchTime = flag.Float64("test.benchtime", 1, "approximate run time for each benchmark, in seconds")
+var benchmarkMemory = flag.Bool("test.benchmem", false, "print memory allocations for benchmarks")
+
+// Global lock to ensure only one benchmark runs at a time.
+var benchmarkLock sync.Mutex
+
+// Used for every benchmark for measuring memory.
+var memStats runtime.MemStats
 
 // An internal type but exported because it is cross-package; part of the implementation
 // of the "go test" command.
@@ -31,6 +39,12 @@ type B struct {
 	bytes     int64
 	timerOn   bool
 	result    BenchmarkResult
+	// The initial states of memStats.Mallocs and memStats.TotalAlloc.
+	startAllocs uint64
+	startBytes  uint64
+	// The net total of this test after being run.
+	netAllocs uint64
+	netBytes  uint64
 }
 
 // StartTimer starts timing a test.  This function is called automatically
@@ -38,6 +52,9 @@ type B struct {
 // a call to StopTimer.
 func (b *B) StartTimer() {
 	if !b.timerOn {
+		runtime.ReadMemStats(&memStats)
+		b.startAllocs = memStats.Mallocs
+		b.startBytes = memStats.TotalAlloc
 		b.start = time.Now()
 		b.timerOn = true
 	}
@@ -49,6 +66,9 @@ func (b *B) StartTimer() {
 func (b *B) StopTimer() {
 	if b.timerOn {
 		b.duration += time.Now().Sub(b.start)
+		runtime.ReadMemStats(&memStats)
+		b.netAllocs += memStats.Mallocs - b.startAllocs
+		b.netBytes += memStats.TotalAlloc - b.startBytes
 		b.timerOn = false
 	}
 }
@@ -57,9 +77,14 @@ func (b *B) StopTimer() {
 // It does not affect whether the timer is running.
 func (b *B) ResetTimer() {
 	if b.timerOn {
+		runtime.ReadMemStats(&memStats)
+		b.startAllocs = memStats.Mallocs
+		b.startBytes = memStats.TotalAlloc
 		b.start = time.Now()
 	}
 	b.duration = 0
+	b.netAllocs = 0
+	b.netBytes = 0
 }
 
 // SetBytes records the number of bytes processed in a single operation.
@@ -75,6 +100,8 @@ func (b *B) nsPerOp() int64 {
 
 // runN runs a single benchmark for the specified number of iterations.
 func (b *B) runN(n int) {
+	benchmarkLock.Lock()
+	defer benchmarkLock.Unlock()
 	// Try to get a comparable environment for each run
 	// by clearing garbage from previous runs.
 	runtime.GC()
@@ -168,14 +195,16 @@ func (b *B) launch() {
 		n = roundUp(n)
 		b.runN(n)
 	}
-	b.result = BenchmarkResult{b.N, b.duration, b.bytes}
+	b.result = BenchmarkResult{b.N, b.duration, b.bytes, b.netAllocs, b.netBytes}
 }
 
 // The results of a benchmark run.
 type BenchmarkResult struct {
-	N     int           // The number of iterations.
-	T     time.Duration // The total time taken.
-	Bytes int64         // Bytes processed in one iteration.
+	N         int           // The number of iterations.
+	T         time.Duration // The total time taken.
+	Bytes     int64         // Bytes processed in one iteration.
+	MemAllocs uint64        // The total number of memory allocations.
+	MemBytes  uint64        // The total number of bytes allocated.
 }
 
 func (r BenchmarkResult) NsPerOp() int64 {
@@ -192,6 +221,20 @@ func (r BenchmarkResult) mbPerSec() float64 {
 	return (float64(r.Bytes) * float64(r.N) / 1e6) / r.T.Seconds()
 }
 
+func (r BenchmarkResult) AllocsPerOp() int64 {
+	if r.N <= 0 {
+		return 0
+	}
+	return int64(r.MemAllocs) / int64(r.N)
+}
+
+func (r BenchmarkResult) AllocedBytesPerOp() int64 {
+	if r.N <= 0 {
+		return 0
+	}
+	return int64(r.MemBytes) / int64(r.N)
+}
+
 func (r BenchmarkResult) String() string {
 	mbs := r.mbPerSec()
 	mb := ""
@@ -212,6 +255,11 @@ func (r BenchmarkResult) String() string {
 	return fmt.Sprintf("%8d\t%s%s", r.N, ns, mb)
 }
 
+func (r BenchmarkResult) MemString() string {
+	return fmt.Sprintf("\t%8d B/op\t%8d allocs/op",
+		r.AllocedBytesPerOp(), r.AllocsPerOp())
+}
+
 // An internal function but exported because it is cross-package; part of the implementation
 // of the "go test" command.
 func RunBenchmarks(matchString func(pat, str string) (bool, error), benchmarks []InternalBenchmark) {
@@ -249,7 +297,11 @@ func RunBenchmarks(matchString func(pat, str string) (bool, error), benchmarks [
 				fmt.Printf("--- FAIL: %s\n%s", benchName, b.output)
 				continue
 			}
-			fmt.Printf("%v\n", r)
+			results := r.String()
+			if *benchmarkMemory {
+				results += "\t" + r.MemString()
+			}
+			fmt.Println(results)
 			// Unlike with tests, we ignore the -chatty flag and always print output for
 			// benchmarks since the output generation time will skew the results.
 			if len(b.output) > 0 {

src/pkg/testing/testing.go

@@ -13,7 +13,7 @@
 // Functions of the form
 //     func BenchmarkXxx(*testing.B)
 // are considered benchmarks, and are executed by the "go test" command when
-// the -test.bench flag is provided.
+// the -test.bench flag is provided. Benchmarks are run sequentially.
 //
 // A sample benchmark function looks like this:
 //     func BenchmarkHello(b *testing.B) {