testing: add AllocsPerRun

This CL also replaces similar loops in other stdlib
package tests with calls to AllocsPerRun.

Fixes #4461.

R=minux.ma, rsc
CC=golang-dev
https://golang.org/cl/7002055

src/pkg/encoding/gob/timing_test.go

@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"runtime"
 	"testing"
 )
 
@@ -50,49 +49,43 @@ func BenchmarkEndToEndByteBuffer(b *testing.B) {
 }
 
 func TestCountEncodeMallocs(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
+	const N = 1000
+
 	var buf bytes.Buffer
 	enc := NewEncoder(&buf)
 	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}
-	memstats := new(runtime.MemStats)
-	runtime.ReadMemStats(memstats)
-	mallocs := 0 - memstats.Mallocs
-	const count = 1000
-	for i := 0; i < count; i++ {
+
+	allocs := testing.AllocsPerRun(N, func() {
 		err := enc.Encode(bench)
 		if err != nil {
 			t.Fatal("encode:", err)
 		}
-	}
-	runtime.ReadMemStats(memstats)
-	mallocs += memstats.Mallocs
-	fmt.Printf("mallocs per encode of type Bench: %d\n", mallocs/count)
+	})
+	fmt.Printf("mallocs per encode of type Bench: %v\n", allocs)
 }
 
 func TestCountDecodeMallocs(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
+	const N = 1000
+
 	var buf bytes.Buffer
 	enc := NewEncoder(&buf)
 	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}
-	const count = 1000
-	for i := 0; i < count; i++ {
+
+	// Fill the buffer with enough to decode
+	testing.AllocsPerRun(N, func() {
 		err := enc.Encode(bench)
 		if err != nil {
 			t.Fatal("encode:", err)
 		}
-	}
+	})
+
 	dec := NewDecoder(&buf)
-	memstats := new(runtime.MemStats)
-	runtime.ReadMemStats(memstats)
-	mallocs := 0 - memstats.Mallocs
-	for i := 0; i < count; i++ {
+	allocs := testing.AllocsPerRun(N, func() {
 		*bench = Bench{}
 		err := dec.Decode(&bench)
 		if err != nil {
 			t.Fatal("decode:", err)
 		}
-	}
-	runtime.ReadMemStats(memstats)
-	mallocs += memstats.Mallocs
-	fmt.Printf("mallocs per decode of type Bench: %d\n", mallocs/count)
+	})
+	fmt.Printf("mallocs per decode of type Bench: %v\n", allocs)
 }

src/pkg/fmt/fmt_test.go

@@ -9,7 +9,6 @@ import (
 	. "fmt"
 	"io"
 	"math"
-	"runtime" // for the malloc count test only
 	"strings"
 	"testing"
 	"time"
@@ -598,19 +597,10 @@ var mallocTest = []struct {
 var _ bytes.Buffer
 
 func TestCountMallocs(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
 	for _, mt := range mallocTest {
-		const N = 100
-		memstats := new(runtime.MemStats)
-		runtime.ReadMemStats(memstats)
-		mallocs := 0 - memstats.Mallocs
-		for i := 0; i < N; i++ {
-			mt.fn()
-		}
-		runtime.ReadMemStats(memstats)
-		mallocs += memstats.Mallocs
-		if mallocs/N > uint64(mt.count) {
-			t.Errorf("%s: expected %d mallocs, got %d", mt.desc, mt.count, mallocs/N)
+		mallocs := testing.AllocsPerRun(100, mt.fn)
+		if got, max := mallocs, float64(mt.count); got > max {
+			t.Errorf("%s: got %v allocs, want <=%v", mt.desc, got, max)
 		}
 	}
 }

src/pkg/net/http/header_test.go

@@ -6,7 +6,6 @@ package http
 
 import (
 	"bytes"
-	"runtime"
 	"testing"
 	"time"
 )
@@ -175,38 +174,31 @@ func TestHasToken(t *testing.T) {
 	}
 }
 
-func BenchmarkHeaderWriteSubset(b *testing.B) {
-	doHeaderWriteSubset(b.N, b)
+var testHeader = Header{
+	"Content-Length": {"123"},
+	"Content-Type":   {"text/plain"},
+	"Date":           {"some date at some time Z"},
+	"Server":         {"Go http package"},
 }
 
-func TestHeaderWriteSubsetMallocs(t *testing.T) {
-	doHeaderWriteSubset(100, t)
-}
+var buf bytes.Buffer
 
-type errorfer interface {
-	Errorf(string, ...interface{})
+func BenchmarkHeaderWriteSubset(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		buf.Reset()
+		testHeader.WriteSubset(&buf, nil)
+	}
 }
 
-func doHeaderWriteSubset(n int, t errorfer) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
-	h := Header(map[string][]string{
-		"Content-Length": {"123"},
-		"Content-Type":   {"text/plain"},
-		"Date":           {"some date at some time Z"},
-		"Server":         {"Go http package"},
-	})
-	var buf bytes.Buffer
-	var m0 runtime.MemStats
-	runtime.ReadMemStats(&m0)
-	for i := 0; i < n; i++ {
+func TestHeaderWriteSubsetMallocs(t *testing.T) {
+	n := testing.AllocsPerRun(100, func() {
 		buf.Reset()
-		h.WriteSubset(&buf, nil)
-	}
-	var m1 runtime.MemStats
-	runtime.ReadMemStats(&m1)
-	if mallocs := m1.Mallocs - m0.Mallocs; n >= 100 && mallocs >= uint64(n) {
+		testHeader.WriteSubset(&buf, nil)
+	})
+	if n > 1 {
 		// TODO(bradfitz,rsc): once we can sort without allocating,
 		// make this an error.  See http://golang.org/issue/3761
-		// t.Errorf("did %d mallocs (>= %d iterations); should have avoided mallocs", mallocs, n)
+		// t.Errorf("got %v allocs, want <= %v", n, 1)
 	}
 }

src/pkg/net/rpc/server_test.go

@@ -445,8 +445,7 @@ func dialHTTP() (*Client, error) {
 	return DialHTTP("tcp", httpServerAddr)
 }
 
-func countMallocs(dial func() (*Client, error), t *testing.T) uint64 {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
+func countMallocs(dial func() (*Client, error), t *testing.T) float64 {
 	once.Do(startServer)
 	client, err := dial()
 	if err != nil {
@@ -454,11 +453,7 @@ func countMallocs(dial func() (*Client, error), t *testing.T) uint64 {
 	}
 	args := &Args{7, 8}
 	reply := new(Reply)
-	memstats := new(runtime.MemStats)
-	runtime.ReadMemStats(memstats)
-	mallocs := 0 - memstats.Mallocs
-	const count = 100
-	for i := 0; i < count; i++ {
+	return testing.AllocsPerRun(100, func() {
 		err := client.Call("Arith.Add", args, reply)
 		if err != nil {
 			t.Errorf("Add: expected no error but got string %q", err.Error())
@@ -466,18 +461,15 @@ func countMallocs(dial func() (*Client, error), t *testing.T) uint64 {
 		if reply.C != args.A+args.B {
 			t.Errorf("Add: expected %d got %d", reply.C, args.A+args.B)
 		}
-	}
-	runtime.ReadMemStats(memstats)
-	mallocs += memstats.Mallocs
-	return mallocs / count
+	})
 }
 
 func TestCountMallocs(t *testing.T) {
-	fmt.Printf("mallocs per rpc round trip: %d\n", countMallocs(dialDirect, t))
+	fmt.Printf("mallocs per rpc round trip: %v\n", countMallocs(dialDirect, t))
 }
 
 func TestCountMallocsOverHTTP(t *testing.T) {
-	fmt.Printf("mallocs per HTTP rpc round trip: %d\n", countMallocs(dialHTTP, t))
+	fmt.Printf("mallocs per HTTP rpc round trip: %v\n", countMallocs(dialHTTP, t))
 }
 
 type writeCrasher struct {

src/pkg/path/filepath/path_test.go

@@ -91,7 +91,6 @@ var wincleantests = []PathTest{
 }
 
 func TestClean(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
 	tests := cleantests
 	if runtime.GOOS == "windows" {
 		for i := range tests {
@@ -108,20 +107,12 @@ func TestClean(t *testing.T) {
 		}
 	}
 
-	var ms runtime.MemStats
-	runtime.ReadMemStats(&ms)
-	allocs := -ms.Mallocs
-	const rounds = 100
-	for i := 0; i < rounds; i++ {
-		for _, test := range tests {
-			filepath.Clean(test.result)
+	for _, test := range tests {
+		allocs := testing.AllocsPerRun(100, func() { filepath.Clean(test.result) })
+		if allocs > 0 {
+			t.Errorf("Clean(%q): %v allocs, want zero", test.result, allocs)
 		}
 	}
-	runtime.ReadMemStats(&ms)
-	allocs += ms.Mallocs
-	if allocs >= rounds {
-		t.Errorf("Clean cleaned paths: %d allocations per test round, want zero", allocs/rounds)
-	}
 }
 
 const sep = filepath.Separator

src/pkg/path/path_test.go

@@ -5,7 +5,6 @@
 package path
 
 import (
-	"runtime"
 	"testing"
 )
 
@@ -64,7 +63,6 @@ var cleantests = []PathTest{
 }
 
 func TestClean(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
 	for _, test := range cleantests {
 		if s := Clean(test.path); s != test.result {
 			t.Errorf("Clean(%q) = %q, want %q", test.path, s, test.result)
@@ -74,20 +72,12 @@ func TestClean(t *testing.T) {
 		}
 	}
 
-	var ms runtime.MemStats
-	runtime.ReadMemStats(&ms)
-	allocs := -ms.Mallocs
-	const rounds = 100
-	for i := 0; i < rounds; i++ {
-		for _, test := range cleantests {
-			Clean(test.result)
+	for _, test := range cleantests {
+		allocs := testing.AllocsPerRun(100, func() { Clean(test.result) })
+		if allocs > 0 {
+			t.Errorf("Clean(%q): %v allocs, want zero", test.result, allocs)
 		}
 	}
-	runtime.ReadMemStats(&ms)
-	allocs += ms.Mallocs
-	if allocs >= rounds {
-		t.Errorf("Clean cleaned paths: %d allocations per test round, want zero", allocs/rounds)
-	}
 }
 
 type SplitTest struct {

src/pkg/reflect/all_test.go

@@ -13,7 +13,6 @@ import (
 	"math/rand"
 	"os"
 	. "reflect"
-	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -2012,20 +2011,13 @@ func TestAddr(t *testing.T) {
 }
 
 func noAlloc(t *testing.T, n int, f func(int)) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
-	// once to prime everything
-	f(-1)
-	memstats := new(runtime.MemStats)
-	runtime.ReadMemStats(memstats)
-	oldmallocs := memstats.Mallocs
-
-	for j := 0; j < n; j++ {
-		f(j)
-	}
-	runtime.ReadMemStats(memstats)
-	mallocs := memstats.Mallocs - oldmallocs
-	if mallocs > 0 {
-		t.Fatalf("%d mallocs after %d iterations", mallocs, n)
+	i := -1
+	allocs := testing.AllocsPerRun(n, func() {
+		f(i)
+		i++
+	})
+	if allocs > 0 {
+		t.Errorf("%d iterations: got %v mallocs, want 0", n, allocs)
 	}
 }
 

src/pkg/strconv/strconv_test.go

@@ -5,7 +5,6 @@
 package strconv_test
 
 import (
-	"runtime"
 	. "strconv"
 	"strings"
 	"testing"
@@ -44,19 +43,10 @@ var (
 )
 
 func TestCountMallocs(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
 	for _, mt := range mallocTest {
-		const N = 100
-		memstats := new(runtime.MemStats)
-		runtime.ReadMemStats(memstats)
-		mallocs := 0 - memstats.Mallocs
-		for i := 0; i < N; i++ {
-			mt.fn()
-		}
-		runtime.ReadMemStats(memstats)
-		mallocs += memstats.Mallocs
-		if mallocs/N > uint64(mt.count) {
-			t.Errorf("%s: expected %d mallocs, got %d", mt.desc, mt.count, mallocs/N)
+		allocs := testing.AllocsPerRun(100, mt.fn)
+		if max := float64(mt.count); allocs > max {
+			t.Errorf("%s: %v allocs, want <=%v", mt.desc, allocs, max)
 		}
 	}
 }

src/pkg/testing/allocs.go

+// Copyright 2013 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 testing
+
+import (
+	"runtime"
+)
+
+// AllocsPerRun returns the average number of allocations during calls to f.
+//
+// To compute the number of allocations, the function will first be run once as
+// a warm-up.  The average number of allocations over the specified number of
+// runs will then be measured and returned.
+//
+// AllocsPerRun sets GOMAXPROCS to 1 during its measurement and will restore
+// it before returning.
+func AllocsPerRun(runs int, f func()) (avg float64) {
+	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
+
+	// Warm up the function
+	f()
+
+	// Measure the starting statistics
+	var memstats runtime.MemStats
+	runtime.ReadMemStats(&memstats)
+	mallocs := 0 - memstats.Mallocs
+
+	// Run the function the specified number of times
+	for i := 0; i < runs; i++ {
+		f()
+	}
+
+	// Read the final statistics
+	runtime.ReadMemStats(&memstats)
+	mallocs += memstats.Mallocs
+
+	// Average the mallocs over the runs (not counting the warm-up)
+	return float64(mallocs) / float64(runs)
+}

src/pkg/time/time_test.go

@@ -11,7 +11,6 @@ import (
 	"fmt"
 	"math/big"
 	"math/rand"
-	"runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -1258,19 +1257,10 @@ var mallocTest = []struct {
 }
 
 func TestCountMallocs(t *testing.T) {
-	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(1))
 	for _, mt := range mallocTest {
-		const N = 100
-		memstats := new(runtime.MemStats)
-		runtime.ReadMemStats(memstats)
-		mallocs := 0 - memstats.Mallocs
-		for i := 0; i < N; i++ {
-			mt.fn()
-		}
-		runtime.ReadMemStats(memstats)
-		mallocs += memstats.Mallocs
-		if mallocs/N > uint64(mt.count) {
-			t.Errorf("%s: expected %d mallocs, got %d", mt.desc, mt.count, mallocs/N)
+		allocs := int(testing.AllocsPerRun(100, mt.fn))
+		if allocs > mt.count {
+			t.Errorf("%s: %d allocs, want %d", mt.desc, allocs, mt.count)
 		}
 	}
 }