testing: add TB.Helper to better support test helpers

This CL implements the proposal at
https://github.com/golang/proposal/blob/master/design/4899-testing-helper.md.

It's based on Josh's CL 79890043 from a few years ago:
https://codereview.appspot.com/79890043 but makes several changes,
most notably by using the new CallersFrames API so that it works with
mid-stack inlining.

Another detail came up while I was working on this: I didn't want the
user to be able to call t.Helper from inside their TestXxx function
directly (which would mean we'd print a file:line from inside the
testing package itself), so I explicitly prevented this from working.

Fixes #4899.

Change-Id: I37493edcfb63307f950442bbaf993d1589515310
Reviewed-on: https://go-review.googlesource.com/38796
Run-TryBot: Caleb Spare <cespare@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

1 files changed, 85 insertions, 14 deletions

diff --git a/src/testing/testing.go b/src/testing/testing.go
index e5a3c3b93c..f95f4ec4a6 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -273,17 +273,20 @@ var (
 // common holds the elements common between T and B and
 // captures common methods such as Errorf.
 type common struct {
-	mu         sync.RWMutex // guards output, w, failed, and done.
-	output     []byte       // Output generated by test or benchmark.
-	w          io.Writer    // For flushToParent.
-	chatty     bool         // A copy of the chatty flag.
-	ran        bool         // Test or benchmark (or one of its subtests) was executed.
-	failed     bool         // Test or benchmark has failed.
-	skipped    bool         // Test of benchmark has been skipped.
-	finished   bool         // Test function has completed.
-	done       bool         // Test is finished and all subtests have completed.
-	hasSub     int32        // written atomically
-	raceErrors int          // number of races detected during test
+	mu      sync.RWMutex         // guards this group of fields
+	output  []byte               // Output generated by test or benchmark.
+	w       io.Writer            // For flushToParent.
+	ran     bool                 // Test or benchmark (or one of its subtests) was executed.
+	failed  bool                 // Test or benchmark has failed.
+	skipped bool                 // Test of benchmark has been skipped.
+	done    bool                 // Test is finished and all subtests have completed.
+	helpers map[uintptr]struct{} // functions to be skipped when writing file/line info
+
+	chatty     bool    // A copy of the chatty flag.
+	finished   bool    // Test function has completed.
+	hasSub     int32   // written atomically
+	raceErrors int     // number of races detected during test
+	runner     uintptr // entry pc of tRunner running the test
 
 	parent   *common
 	level    int       // Nesting depth of test or benchmark.
@@ -312,10 +315,48 @@ func Verbose() bool {
 	return *chatty
 }
 
+// frameSkip searches, starting after skip frames, for the first caller frame
+// in a function not marked as a helper and returns the frames to skip
+// to reach that site. The search stops if it finds a tRunner function that
+// was the entry point into the test.
+// This function must be called with c.mu held.
+func (c *common) frameSkip(skip int) int {
+	if c.helpers == nil {
+		return skip
+	}
+	var pc [50]uintptr
+	// Skip two extra frames to account for this function
+	// and runtime.Callers itself.
+	n := runtime.Callers(skip+2, pc[:])
+	if n == 0 {
+		panic("testing: zero callers found")
+	}
+	frames := runtime.CallersFrames(pc[:n])
+	var frame runtime.Frame
+	more := true
+	for i := 0; more; i++ {
+		frame, more = frames.Next()
+		if frame.Entry == c.runner {
+			// We've gone up all the way to the tRunner calling
+			// the test function (so the user must have
+			// called tb.Helper from inside that test function).
+			// Only skip up to the test function itself.
+			return skip + i - 1
+		}
+		if _, ok := c.helpers[frame.Entry]; !ok {
+			// Found a frame that wasn't inside a helper function.
+			return skip + i
+		}
+	}
+	return skip
+}
+
 // decorate prefixes the string with the file and line of the call site
 // and inserts the final newline if needed and indentation tabs for formatting.
-func decorate(s string) string {
-	_, file, line, ok := runtime.Caller(3) // decorate + log + public function.
+// This function must be called with c.mu held.
+func (c *common) decorate(s string) string {
+	skip := c.frameSkip(3) // decorate + log + public function.
+	_, file, line, ok := runtime.Caller(skip)
 	if ok {
 		// Truncate file name at last file name separator.
 		if index := strings.LastIndex(file, "/"); index >= 0 {
@@ -405,6 +446,7 @@ type TB interface {
 	SkipNow()
 	Skipf(format string, args ...interface{})
 	Skipped() bool
+	Helper()
 
 	// A private method to prevent users implementing the
 	// interface and so future additions to it will not
@@ -505,7 +547,7 @@ func (c *common) FailNow() {
 func (c *common) log(s string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	c.output = append(c.output, decorate(s)...)
+	c.output = append(c.output, c.decorate(s)...)
 }
 
 // Log formats its arguments using default formatting, analogous to Println,
@@ -583,6 +625,33 @@ func (c *common) Skipped() bool {
 	return c.skipped
 }
 
+// Helper marks the calling function as a test helper function.
+// When printing file and line information, that function will be skipped.
+// Helper may be called simultaneously from multiple goroutines.
+// Helper has no effect if it is called directly from a TestXxx/BenchmarkXxx
+// function or a subtest/sub-benchmark function.
+func (c *common) Helper() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.helpers == nil {
+		c.helpers = make(map[uintptr]struct{})
+	}
+	c.helpers[callerEntry(1)] = struct{}{}
+}
+
+// callerEntry gives the entry pc for the caller after skip frames
+// (where 0 means the current function).
+func callerEntry(skip int) uintptr {
+	var pc [1]uintptr
+	n := runtime.Callers(skip+2, pc[:]) // skip + runtime.Callers + callerEntry
+	if n == 0 {
+		panic("testing: zero callers found")
+	}
+	frames := runtime.CallersFrames(pc[:])
+	frame, _ := frames.Next()
+	return frame.Entry
+}
+
 // Parallel signals that this test is to be run in parallel with (and only with)
 // other parallel tests. When a test is run multiple times due to use of
 // -test.count or -test.cpu, multiple instances of a single test never run in
@@ -617,6 +686,8 @@ type InternalTest struct {
 }
 
 func tRunner(t *T, fn func(t *T)) {
+	t.runner = callerEntry(0)
+
 	// When this goroutine is done, either because fn(t)
 	// returned normally or because a test failure triggered
 	// a call to runtime.Goexit, record the duration and send