2 files changed, 70 insertions, 1 deletions

diff --git a/src/testing/benchmark.go b/src/testing/benchmark.go
index 0e31ef0e66..bc9be9b176 100644
--- a/src/testing/benchmark.go
+++ b/src/testing/benchmark.go
@@ -124,6 +124,8 @@ type B struct {
 		// i is the current Loop iteration. It's strictly monotonically
 		// increasing toward n.
 		i int
+
+		done bool // set when B.Loop return false
 	}
 }
 
@@ -201,6 +203,7 @@ func (b *B) runN(n int) {
 	b.N = n
 	b.loop.n = 0
 	b.loop.i = 0
+	b.loop.done = false
 	b.ctx = ctx
 	b.cancelCtx = cancelCtx
 
@@ -211,6 +214,10 @@ func (b *B) runN(n int) {
 	b.StopTimer()
 	b.previousN = n
 	b.previousDuration = b.duration
+
+	if b.loop.n > 0 && !b.loop.done && !b.failed {
+		b.Error("benchmark function returned without B.Loop() == false (break or return in loop?)")
+	}
 }
 
 // run1 runs the first iteration of benchFunc. It reports whether more
@@ -382,6 +389,7 @@ func (b *B) stopOrScaleBLoop() bool {
 		b.StopTimer()
 		// Commit iteration count
 		b.N = b.loop.n
+		b.loop.done = true
 		return false
 	}
 	// Loop scaling
@@ -414,6 +422,7 @@ func (b *B) loopSlowPath() bool {
 		b.StopTimer()
 		// Commit iteration count
 		b.N = b.loop.n
+		b.loop.done = true
 		return false
 	}
 	// Handles fixed time case
diff --git a/src/testing/loop_test.go b/src/testing/loop_test.go
index 7a42919643..423094fbbd 100644
--- a/src/testing/loop_test.go
+++ b/src/testing/loop_test.go
@@ -4,6 +4,13 @@
 
 package testing
 
+import (
+	"bytes"
+	"strings"
+)
+
+// See also TestBenchmarkBLoop* in other files.
+
 func TestBenchmarkBLoop(t *T) {
 	var initialStart highPrecisionTime
 	var firstStart highPrecisionTime
@@ -68,4 +75,57 @@ func TestBenchmarkBLoop(t *T) {
 	}
 }
 
-// See also TestBenchmarkBLoop* in other files.
+func TestBenchmarkBLoopBreak(t *T) {
+	var bState *B
+	var bLog bytes.Buffer
+	bRet := Benchmark(func(b *B) {
+		// The Benchmark function provides no access to the failure state and
+		// discards the log, so capture the B and save its log.
+		bState = b
+		b.common.w = &bLog
+
+		for i := 0; b.Loop(); i++ {
+			if i == 2 {
+				break
+			}
+		}
+	})
+	if !bState.failed {
+		t.Errorf("benchmark should have failed")
+	}
+	const wantLog = "benchmark function returned without B.Loop"
+	if log := bLog.String(); !strings.Contains(log, wantLog) {
+		t.Errorf("missing error %q in output:\n%s", wantLog, log)
+	}
+	// A benchmark that exits early should not report its target iteration count
+	// because it's not meaningful.
+	if bRet.N != 0 {
+		t.Errorf("want N == 0, got %d", bRet.N)
+	}
+}
+
+func TestBenchmarkBLoopError(t *T) {
+	// Test that a benchmark that exits early because of an error doesn't *also*
+	// complain that the benchmark exited early.
+	var bState *B
+	var bLog bytes.Buffer
+	bRet := Benchmark(func(b *B) {
+		bState = b
+		b.common.w = &bLog
+
+		for i := 0; b.Loop(); i++ {
+			b.Error("error")
+			return
+		}
+	})
+	if !bState.failed {
+		t.Errorf("benchmark should have failed")
+	}
+	const noWantLog = "benchmark function returned without B.Loop"
+	if log := bLog.String(); strings.Contains(log, noWantLog) {
+		t.Errorf("unexpected error %q in output:\n%s", noWantLog, log)
+	}
+	if bRet.N != 0 {
+		t.Errorf("want N == 0, got %d", bRet.N)
+	}
+}