Merge pull request #1614 from DevotedHealth/mauclair-call-stack-perf

Call stack perf change for CallerInfo

Author: brackendawson

assert/assertions.go

@@ -24,6 +24,8 @@ import (
 	"github.com/stretchr/testify/assert/yaml"
 )
 
+const stackFrameBufferSize = 10
+
 //go:generate sh -c "cd ../_codegen && go build && cd - && ../_codegen/_codegen -output-package=assert -template=assertion_format.go.tmpl"
 
 // TestingT is an interface wrapper around *testing.T
@@ -212,57 +214,73 @@ the problem actually occurred in calling code.*/
 func CallerInfo() []string {
 
 	var pc uintptr
-	var ok bool
 	var file string
 	var line int
 	var name string
 
 	callers := []string{}
-	for i := 0; ; i++ {
-		pc, file, line, ok = runtime.Caller(i)
-		if !ok {
-			// The breaks below failed to terminate the loop, and we ran off the
-			// end of the call stack.
-			break
-		}
+	pcs := make([]uintptr, stackFrameBufferSize)
+	offset := 1
 
-		// This is a huge edge case, but it will panic if this is the case, see #180
-		if file == "<autogenerated>" {
-			break
-		}
+	for {
+		n := runtime.Callers(offset, pcs)
 
-		f := runtime.FuncForPC(pc)
-		if f == nil {
-			break
-		}
-		name = f.Name()
-
-		// testing.tRunner is the standard library function that calls
-		// tests. Subtests are called directly by tRunner, without going through
-		// the Test/Benchmark/Example function that contains the t.Run calls, so
-		// with subtests we should break when we hit tRunner, without adding it
-		// to the list of callers.
-		if name == "testing.tRunner" {
+		if n == 0 {
 			break
 		}
 
-		parts := strings.Split(file, "/")
-		if len(parts) > 1 {
-			filename := parts[len(parts)-1]
-			dir := parts[len(parts)-2]
-			if (dir != "assert" && dir != "mock" && dir != "require") || filename == "mock_test.go" {
-				callers = append(callers, fmt.Sprintf("%s:%d", file, line))
+		frames := runtime.CallersFrames(pcs[:n])
+
+		for {
+			frame, more := frames.Next()
+			pc = frame.PC
+			file = frame.File
+			line = frame.Line
+
+			// This is a huge edge case, but it will panic if this is the case, see #180
+			if file == "<autogenerated>" {
+				break
 			}
-		}
 
-		// Drop the package
-		segments := strings.Split(name, ".")
-		name = segments[len(segments)-1]
-		if isTest(name, "Test") ||
-			isTest(name, "Benchmark") ||
-			isTest(name, "Example") {
-			break
+			f := runtime.FuncForPC(pc)
+			if f == nil {
+				break
+			}
+			name = f.Name()
+
+			// testing.tRunner is the standard library function that calls
+			// tests. Subtests are called directly by tRunner, without going through
+			// the Test/Benchmark/Example function that contains the t.Run calls, so
+			// with subtests we should break when we hit tRunner, without adding it
+			// to the list of callers.
+			if name == "testing.tRunner" {
+				break
+			}
+
+			parts := strings.Split(file, "/")
+			if len(parts) > 1 {
+				filename := parts[len(parts)-1]
+				dir := parts[len(parts)-2]
+				if (dir != "assert" && dir != "mock" && dir != "require") || filename == "mock_test.go" {
+					callers = append(callers, fmt.Sprintf("%s:%d", file, line))
+				}
+			}
+
+			// Drop the package
+			segments := strings.Split(name, ".")
+			name = segments[len(segments)-1]
+			if isTest(name, "Test") ||
+				isTest(name, "Benchmark") ||
+				isTest(name, "Example") {
+				break
+			}
+			if !more {
+				break
+			}
 		}
+		// We know we already have less than a buffer's worth of frames
+		offset += stackFrameBufferSize
+
 	}
 
 	return callers