testing: make Cleanup work for benchmarks too.

Fixes #37073.

Change-Id: I6fb24a3f9d7b7adf3213ac6a8bcbf5fb43975b7e
Reviewed-on: https://go-review.googlesource.com/c/go/+/218117
Reviewed-by: Bryan C. Mills <bcmills@google.com>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

2 files changed, 41 insertions, 0 deletions

diff --git a/src/testing/benchmark.go b/src/testing/benchmark.go
index 93f461b07a..88ba0f0242 100644
--- a/src/testing/benchmark.go
+++ b/src/testing/benchmark.go
@@ -179,6 +179,7 @@ func (b *B) ReportAllocs() {
 func (b *B) runN(n int) {
 	benchmarkLock.Lock()
 	defer benchmarkLock.Unlock()
+	defer b.runCleanup(normalPanic)
 	// Try to get a comparable environment for each run
 	// by clearing garbage from previous runs.
 	runtime.GC()
diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
index 3dc30ee72e..95f8220f81 100644
--- a/src/testing/sub_test.go
+++ b/src/testing/sub_test.go
@@ -613,6 +613,46 @@ func TestBRun(t *T) {
 				t.Errorf("MemBytes was %v; want %v", got, 2*bufSize)
 			}
 		},
+	}, {
+		desc: "cleanup is called",
+		f: func(b *B) {
+			var calls, cleanups, innerCalls, innerCleanups int
+			b.Run("", func(b *B) {
+				calls++
+				b.Cleanup(func() {
+					cleanups++
+				})
+				b.Run("", func(b *B) {
+					b.Cleanup(func() {
+						innerCleanups++
+					})
+					innerCalls++
+				})
+				work(b)
+			})
+			if calls == 0 || calls != cleanups {
+				t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+			}
+			if innerCalls == 0 || innerCalls != innerCleanups {
+				t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+			}
+		},
+	}, {
+		desc:   "cleanup is called on failure",
+		failed: true,
+		f: func(b *B) {
+			var calls, cleanups int
+			b.Run("", func(b *B) {
+				calls++
+				b.Cleanup(func() {
+					cleanups++
+				})
+				b.Fatalf("failure")
+			})
+			if calls == 0 || calls != cleanups {
+				t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+			}
+		},
 	}}
 	for _, tc := range testCases {
 		var ok bool