1 files changed, 57 insertions, 37 deletions

diff --git a/src/testing/benchmark.go b/src/testing/benchmark.go
index 9491213ef1..da20c1c0c8 100644
--- a/src/testing/benchmark.go
+++ b/src/testing/benchmark.go
@@ -176,6 +176,10 @@ func (b *B) ReportAllocs() {
 
 // runN runs a single benchmark for the specified number of iterations.
 func (b *B) runN(n int) {
+	if b.done {
+		panic("testing: internal error: runN when benchmark is already done")
+	}
+
 	benchmarkLock.Lock()
 	defer benchmarkLock.Unlock()
 	defer func() {
@@ -196,46 +200,46 @@ func (b *B) runN(n int) {
 	b.previousDuration = b.duration
 }
 
-// run1 runs the first iteration of benchFunc. It reports whether more
-// iterations of this benchmarks should be run.
-func (b *B) run1() bool {
+// run1 runs the first iteration of benchFunc.
+//
+// If no more iterations of this benchmark should be done, run1 sets b.done.
+func (b *B) run1() {
 	if ctx := b.context; ctx != nil {
 		// Extend maxLen, if needed.
 		if n := len(b.name) + ctx.extLen + 1; n > ctx.maxLen {
 			ctx.maxLen = n + 8 // Add additional slack to avoid too many jumps in size.
 		}
 	}
+
+	runOneDone := make(chan struct{})
 	go func() {
 		// Signal that we're done whether we return normally
 		// or by FailNow's runtime.Goexit.
-		defer func() {
-			b.signal <- true
-		}()
-
+		defer close(runOneDone)
 		b.runN(1)
 	}()
-	<-b.signal
+	<-runOneDone
+
 	if b.failed {
 		fmt.Fprintf(b.w, "%s--- FAIL: %s\n%s", b.chatty.prefix(), b.name, b.output)
-		return false
+		b.done = true
+		close(b.doneOrParallel)
+		return
 	}
 	// Only print the output if we know we are not going to proceed.
 	// Otherwise it is printed in processBench.
-	b.mu.RLock()
-	finished := b.finished
-	b.mu.RUnlock()
-	if b.hasSub.Load() || finished {
+	if b.hasSub.Load() || b.skipped {
 		tag := "BENCH"
 		if b.skipped {
 			tag = "SKIP"
 		}
-		if b.chatty != nil && (len(b.output) > 0 || finished) {
+		if b.chatty != nil && (len(b.output) > 0 || b.skipped) {
 			b.trimOutput()
 			fmt.Fprintf(b.w, "%s--- %s: %s\n%s", b.chatty.prefix(), tag, b.name, b.output)
 		}
-		return false
+		b.done = true
+		close(b.doneOrParallel)
 	}
-	return true
 }
 
 var labelsOnce sync.Once
@@ -262,9 +266,10 @@ func (b *B) run() {
 	}
 }
 
+// doBench calls b.launch in a separate goroutine and waits for it to complete.
 func (b *B) doBench() BenchmarkResult {
 	go b.launch()
-	<-b.signal
+	<-b.doneOrParallel
 	return b.result
 }
 
@@ -276,7 +281,10 @@ func (b *B) launch() {
 	// Signal that we're done whether we return normally
 	// or by FailNow's runtime.Goexit.
 	defer func() {
-		b.signal <- true
+		b.result = BenchmarkResult{b.N, b.duration, b.bytes, b.netAllocs, b.netBytes, b.extra}
+		b.setRanLeaf()
+		b.done = true
+		close(b.doneOrParallel)
 	}()
 
 	// Run the benchmark for at least the specified amount of time.
@@ -316,7 +324,6 @@ func (b *B) launch() {
 			b.runN(int(n))
 		}
 	}
-	b.result = BenchmarkResult{b.N, b.duration, b.bytes, b.netAllocs, b.netBytes, b.extra}
 }
 
 // Elapsed returns the measured elapsed time of the benchmark.
@@ -550,6 +557,7 @@ func runBenchmarks(importPath string, matchString func(pat, str string) (bool, e
 		main.chatty = newChattyPrinter(main.w)
 	}
 	main.runN(1)
+	main.done = true
 	return !main.failed
 }
 
@@ -568,18 +576,21 @@ func (ctx *benchContext) processBench(b *B) {
 			if i > 0 || j > 0 {
 				b = &B{
 					common: common{
-						signal: make(chan bool),
-						name:   b.name,
-						w:      b.w,
-						chatty: b.chatty,
-						bench:  true,
+						doneOrParallel: make(chan struct{}),
+						name:           b.name,
+						w:              b.w,
+						chatty:         b.chatty,
+						bench:          true,
 					},
 					benchFunc: b.benchFunc,
 					benchTime: b.benchTime,
 				}
 				b.run1()
 			}
-			r := b.doBench()
+			var r BenchmarkResult
+			if !b.done {
+				r = b.doBench()
+			}
 			if b.failed {
 				// The output could be very long here, but probably isn't.
 				// We print it all, regardless, because we don't want to trim the reason
@@ -622,6 +633,13 @@ var hideStdoutForTesting = false
 // A subbenchmark is like any other benchmark. A benchmark that calls Run at
 // least once will not be measured itself and will be called once with N=1.
 func (b *B) Run(name string, f func(b *B)) bool {
+	if b.previousN > 0 {
+		// If the benchmark calls Run we will only call it once with N=1.
+		// If it doesn't call Run the first time we try it, it must not
+		// call run on subsequent invocations either.
+		panic("testing: unexpected call to B.Run during iteration")
+	}
+
 	// Since b has subbenchmarks, we will no longer run it as a benchmark itself.
 	// Release the lock and acquire it on exit to ensure locks stay paired.
 	b.hasSub.Store(true)
@@ -639,14 +657,14 @@ func (b *B) Run(name string, f func(b *B)) bool {
 	n := runtime.Callers(2, pc[:])
 	sub := &B{
 		common: common{
-			signal:  make(chan bool),
-			name:    benchName,
-			parent:  &b.common,
-			level:   b.level + 1,
-			creator: pc[:n],
-			w:       b.w,
-			chatty:  b.chatty,
-			bench:   true,
+			doneOrParallel: make(chan struct{}),
+			name:           benchName,
+			parent:         &b.common,
+			level:          b.level + 1,
+			creator:        pc[:n],
+			w:              b.w,
+			chatty:         b.chatty,
+			bench:          true,
 		},
 		importPath: b.importPath,
 		benchFunc:  f,
@@ -679,7 +697,8 @@ func (b *B) Run(name string, f func(b *B)) bool {
 		}
 	}
 
-	if sub.run1() {
+	sub.run1()
+	if !sub.done {
 		sub.run()
 	}
 	b.add(sub.result)
@@ -823,13 +842,14 @@ func (b *B) SetParallelism(p int) {
 func Benchmark(f func(b *B)) BenchmarkResult {
 	b := &B{
 		common: common{
-			signal: make(chan bool),
-			w:      discard{},
+			doneOrParallel: make(chan struct{}),
+			w:              discard{},
 		},
 		benchFunc: f,
 		benchTime: benchTime,
 	}
-	if b.run1() {
+	b.run1()
+	if !b.done {
 		b.run()
 	}
 	return b.result