From 1e095b76229a8faa329dbdee6de246faabc3cf53 Mon Sep 17 00:00:00 2001
From: Shenghou Ma <minux.ma@gmail.com>
Date: Thu, 17 Jan 2013 18:45:49 +0800
Subject: [PATCH] testing: introduce (*B).ReportAllocs() Calling it will show
 memory allocation statistics for that single benchmark (if -test.benchmem is
 not provided)

R=golang-dev, rsc, kevlar, bradfitz
CC=golang-dev
https://golang.org/cl/7027046
---
 src/pkg/exp/html/parse_test.go |  8 +-------
 src/pkg/exp/html/token_test.go |  8 +-------
 src/pkg/testing/benchmark.go   | 20 ++++++++++++++------
 3 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/src/pkg/exp/html/parse_test.go b/src/pkg/exp/html/parse_test.go
index 7cf2ff4163..4896dfb7a0 100644
--- a/src/pkg/exp/html/parse_test.go
+++ b/src/pkg/exp/html/parse_test.go
@@ -382,15 +382,9 @@ func BenchmarkParser(b *testing.B) {
 	}
 	b.SetBytes(int64(len(buf)))
 	runtime.GC()
-	var ms runtime.MemStats
-	runtime.ReadMemStats(&ms)
-	mallocs := ms.Mallocs
+	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		Parse(bytes.NewBuffer(buf))
 	}
-	b.StopTimer()
-	runtime.ReadMemStats(&ms)
-	mallocs = ms.Mallocs - mallocs
-	b.Logf("%d iterations, %d mallocs per iteration\n", b.N, int(mallocs)/b.N)
 }
diff --git a/src/pkg/exp/html/token_test.go b/src/pkg/exp/html/token_test.go
index 63a8bfc483..14e23467f4 100644
--- a/src/pkg/exp/html/token_test.go
+++ b/src/pkg/exp/html/token_test.go
@@ -634,9 +634,7 @@ func benchmarkTokenizer(b *testing.B, level int) {
 	}
 	b.SetBytes(int64(len(buf)))
 	runtime.GC()
-	var ms runtime.MemStats
-	runtime.ReadMemStats(&ms)
-	mallocs := ms.Mallocs
+	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		z := NewTokenizer(bytes.NewBuffer(buf))
@@ -674,10 +672,6 @@ func benchmarkTokenizer(b *testing.B, level int) {
 			}
 		}
 	}
-	b.StopTimer()
-	runtime.ReadMemStats(&ms)
-	mallocs = ms.Mallocs - mallocs
-	b.Logf("%d iterations, %d mallocs per iteration\n", b.N, int(mallocs)/b.N)
 }
 
 func BenchmarkRawLevelTokenizer(b *testing.B)  { benchmarkTokenizer(b, rawLevel) }
diff --git a/src/pkg/testing/benchmark.go b/src/pkg/testing/benchmark.go
index cb92fab50a..25fb2d6191 100644
--- a/src/pkg/testing/benchmark.go
+++ b/src/pkg/testing/benchmark.go
@@ -34,11 +34,12 @@ type InternalBenchmark struct {
 // timing and to specify the number of iterations to run.
 type B struct {
 	common
-	N         int
-	benchmark InternalBenchmark
-	bytes     int64
-	timerOn   bool
-	result    BenchmarkResult
+	N               int
+	benchmark       InternalBenchmark
+	bytes           int64
+	timerOn         bool
+	showAllocResult bool
+	result          BenchmarkResult
 	// The initial states of memStats.Mallocs and memStats.TotalAlloc.
 	startAllocs uint64
 	startBytes  uint64
@@ -91,6 +92,13 @@ func (b *B) ResetTimer() {
 // If this is called, the benchmark will report ns/op and MB/s.
 func (b *B) SetBytes(n int64) { b.bytes = n }
 
+// ReportAllocs enables malloc statistics for this benchmark.
+// It is equivalent to setting -test.benchmem, but it only affects the
+// benchmark function that calls ReportAllocs.
+func (b *B) ReportAllocs() {
+	b.showAllocResult = true
+}
+
 func (b *B) nsPerOp() int64 {
 	if b.N <= 0 {
 		return 0
@@ -298,7 +306,7 @@ func RunBenchmarks(matchString func(pat, str string) (bool, error), benchmarks [
 				continue
 			}
 			results := r.String()
-			if *benchmarkMemory {
+			if *benchmarkMemory || b.showAllocResult {
 				results += "\t" + r.MemString()
 			}
 			fmt.Println(results)
-- 
2.48.1