testing: simplify concurrency and cleanup logic

While investigating #60083, I found a couple of bugs (notably #61034)
that had slipped through code review in part because the concurrency
patterns used in the testing package were too complex for me to fully
reason about. This change adjusts those patterns to be more in line
with current idioms, and to reduce the number of special cases that
depend on details that should be orthogonal. (For example: the details
of how we invoke the Cleanup functions should not depend on whether
the test happened to run any parallel subtests.)

In the process, this change fixes a handful of bugs:

- Concurrent calls to Run (explicitly allowed by TestParallelSub)
  could previously drive the testcontext.running count negative,
  causing the number of running parallel tests to exceed the -parallel
  flag.

- The -failfast flag now takes effect immediately on failure. It no
  longer delays until the test finishes, and no longer misses failures
  during cleanup (fixing #61034).

- If a Cleanup function calls runtime.Goexit (typically via t.FailNow)
  during a panic, Cleanup functions from its parent tests are no
  longer skipped and buffered logs from its parent tests are now
  flushed.

- The time reported for a test with subtests now includes the time spent
  running those subtests, regardless of whether they are parallel.
  (Previously, non-parallel subtests were included but parallel subtests
  were not.)

- Calls to (*B).Run in iterations after the first are now diagnosed
  with a panic. (This diagnoses badly-behaved benchmarks: if Run is
  called during the first iteration, no subsequent iterations are
  supposed to occur.)

Fixes #61034.

Change-Id: I3797f6ef5210a3d2d5d6c2710d3f35c0219b02ea
Cq-Include-Trybots: luci.golang.try:gotip-linux-amd64-longtest,gotip-linux-amd64-longtest-race,gotip-windows-amd64-longtest
Reviewed-on: https://go-review.googlesource.com/c/go/+/506755
Auto-Submit: Bryan Mills <bcmills@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

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