stack: Parse all functions

internal/stack/stacks.go

@@ -34,10 +34,15 @@
 
 // Stack represents a single Goroutine's stack.
 type Stack struct {
-	id            int
-	state         string
+	id    int
+	state string // e.g. 'running', 'chan receive'
+
+	// The first function on the stack.
 	firstFunction string
 
+	// A set of all functions in the stack,
+	allFunctions map[string]struct{}
+
 	// Full, raw stack trace.
 	fullStack string
 }
@@ -62,6 +67,13 @@
 	return s.firstFunction
 }
 
+// HasFunction reports whether the stack has the given function
+// anywhere in it.
+func (s Stack) HasFunction(name string) bool {
+	_, ok := s.allFunctions[name]
+	return ok
+}
+
 func (s Stack) String() string {
 	return fmt.Sprintf(
 		"Goroutine %v in state %v, with %v on top of the stack:\n%s",
@@ -126,9 +138,9 @@
 		firstFunction string
 		fullStack     bytes.Buffer
 	)
+	funcs := make(map[string]struct{})
 	for p.scan.Scan() {
 		line := p.scan.Text()
-
 		if strings.HasPrefix(line, "goroutine ") {
 			// If we see the goroutine header,
 			// it's the end of this stack.
@@ -140,19 +152,74 @@
 		fullStack.WriteString(line)
 		fullStack.WriteByte('\n') // scanner trims the newline
 
-		// The first line after the header is the top of the stack.
-		if firstFunction == "" {
-			firstFunction, err = parseFirstFunc(line)
-			if err != nil {
-				return Stack{}, fmt.Errorf("extract function: %w", err)
+		if len(line) == 0 {
+			// Empty line usually marks the end of the stack
+			// but we don't want to have to rely on that.
+			// Just skip it.
+			continue
+		}
+
+		funcName, creator, err := parseFuncName(line)
+		if err != nil {
+			return Stack{}, fmt.Errorf("parse function: %w", err)
+		}
+		if !creator {
+			// A function is part of a goroutine's stack
+			// only if it's not a "created by" function.
+			//
+			// The creator function is part of a different stack.
+			// We don't care about it right now.
+			funcs[funcName] = struct{}{}
+			if firstFunction == "" {
+				firstFunction = funcName
+			}
+
+		}
+
+		// The function name followed by a line in the form:
+		//
+		//	<tab>example.com/path/to/package/file.go:123 +0x123
+		//
+		// We don't care about the position so we can skip this line.
+		if p.scan.Scan() {
+			// Be defensive:
+			// Skip the line only if it starts with a tab.
+			bs := p.scan.Bytes()
+			if len(bs) > 0 && bs[0] == '\t' {
+				fullStack.Write(bs)
+				fullStack.WriteByte('\n')
+			} else {
+				// Put it back and let the next iteration handle it
+				// if it doesn't start with a tab.
+				p.scan.Unscan()
 			}
 		}
+
+		if creator {
+			// The "created by" line is the last line of the stack.
+			// We can stop parsing now.
+			//
+			// Note that if tracebackancestors=N is set,
+			// there may be more a traceback of the creator function
+			// following the "created by" line,
+			// but it should not be considered part of this stack.
+			// e.g.,
+			//
+			// created by testing.(*T).Run in goroutine 1
+			//         /usr/lib/go/src/testing/testing.go:1648 +0x3ad
+			// [originating from goroutine 1]:
+			// testing.(*T).Run(...)
+			//         /usr/lib/go/src/testing/testing.go:1649 +0x3ad
+			//
+			break
+		}
 	}
 
 	return Stack{
 		id:            id,
 		state:         state,
 		firstFunction: firstFunction,
+		allFunctions:  funcs,
 		fullStack:     fullStack.String(),
 	}, nil
 }
@@ -176,12 +243,35 @@
 	}
 }
 
-func parseFirstFunc(line string) (string, error) {
-	line = strings.TrimSpace(line)
-	if idx := strings.LastIndex(line, "("); idx > 0 {
-		return line[:idx], nil
+// Parses a single function from the given line.
+// The line is in one of these formats:
+//
+//	example.com/path/to/package.funcName(args...)
+//	example.com/path/to/package.(*typeName).funcName(args...)
+//	created by example.com/path/to/package.funcName
+//	created by example.com/path/to/package.funcName in goroutine [...]
+//
+// Also reports whether the line was a "created by" line.
+func parseFuncName(line string) (name string, creator bool, err error) {
+	if after, ok := strings.CutPrefix(line, "created by "); ok {
+		// The function name is the part after "created by "
+		// and before " in goroutine [...]".
+		idx := strings.Index(after, " in goroutine")
+		if idx >= 0 {
+			after = after[:idx]
+		}
+		name = after
+		creator = true
+	} else if idx := strings.LastIndexByte(line, '('); idx >= 0 {
+		// The function name is the part before the last '('.
+		name = line[:idx]
 	}
-	return "", fmt.Errorf("no function found: %q", line)
+
+	if name == "" {
+		return "", false, fmt.Errorf("no function found: %q", line)
+	}
+
+	return name, creator, nil
 }
 
 // parseGoStackHeader parses a stack header that looks like:

internal/stack/stacks_test.go

@@ -68,39 +68,70 @@ func TestAll(t *testing.T) {
 	sort.Sort(byGoroutineID(got))
 
 	assert.Contains(t, got[0].Full(), "testing.(*T).Run")
+	assert.Contains(t, got[0].allFunctions, "testing.(*T).Run")
+
 	assert.Contains(t, got[1].Full(), "TestAll")
+	assert.Contains(t, got[1].allFunctions, "go.uber.org/goleak/internal/stack.TestAll")
+
 	for i := 0; i < 5; i++ {
 		assert.Contains(t, got[2+i].Full(), "stack.waitForDone")
 	}
 }
 
 func TestCurrent(t *testing.T) {
+	const pkgPrefix = "go.uber.org/goleak/internal/stack"
+
 	got := Current()
 	assert.NotZero(t, got.ID(), "Should get non-zero goroutine id")
 	assert.Equal(t, "running", got.State())
 	assert.Equal(t, "go.uber.org/goleak/internal/stack.getStackBuffer", got.FirstFunction())
 
 	wantFrames := []string{
-		"stack.getStackBuffer",
-		"stack.getStacks",
-		"stack.Current",
-		"stack.Current",
-		"stack.TestCurrent",
+		"getStackBuffer",
+		"getStacks",
+		"Current",
+		"Current",
+		"TestCurrent",
 	}
 	all := got.Full()
 	for _, frame := range wantFrames {
-		assert.Contains(t, all, frame)
+		name := pkgPrefix + "." + frame
+		assert.Contains(t, all, name)
+		assert.True(t, got.HasFunction(name), "missing in stack: %v\n%s", name, all)
 	}
 	assert.Contains(t, got.String(), "in state")
 	assert.Contains(t, got.String(), "on top of the stack")
 
+	assert.Contains(t, all, "stack/stacks_test.go",
+		"file name missing in stack:\n%s", all)
+
 	// Ensure that we are not returning the buffer without slicing it
 	// from getStackBuffer.
 	if len(got.Full()) > 1024 {
 		t.Fatalf("Returned stack is too large")
 	}
 }
 
+func TestCurrentCreatedBy(t *testing.T) {
+	var stack Stack
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		stack = Current()
+	}()
+	<-done
+
+	// The test function created the goroutine
+	// so it won't be part of the stack.
+	assert.False(t, stack.HasFunction("go.uber.org/goleak/internal/stack.TestCurrentCreatedBy"),
+		"TestCurrentCreatedBy should not be in stack:\n%s", stack.Full())
+
+	// However, the nested function should be.
+	assert.True(t,
+		stack.HasFunction("go.uber.org/goleak/internal/stack.TestCurrentCreatedBy.func1"),
+		"TestCurrentCreatedBy.func1 is not in stack:\n%s", stack.Full())
+}
+
 func TestAllLargeStack(t *testing.T) {
 	const (
 		stackDepth    = 100
@@ -134,6 +165,47 @@ func TestAllLargeStack(t *testing.T) {
 	close(done)
 }
 
+func TestParseFuncName(t *testing.T) {
+	tests := []struct {
+		name    string
+		give    string
+		want    string
+		creator bool
+	}{
+		{
+			name: "function",
+			give: "example.com/foo/bar.baz()",
+			want: "example.com/foo/bar.baz",
+		},
+		{
+			name: "method",
+			give: "example.com/foo/bar.(*baz).qux()",
+			want: "example.com/foo/bar.(*baz).qux",
+		},
+		{
+			name:    "created by", // Go 1.20
+			give:    "created by example.com/foo/bar.baz",
+			want:    "example.com/foo/bar.baz",
+			creator: true,
+		},
+		{
+			name:    "created by/in goroutine", // Go 1.21
+			give:    "created by example.com/foo/bar.baz in goroutine 123",
+			want:    "example.com/foo/bar.baz",
+			creator: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, creator, err := parseFuncName(tt.give)
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+			assert.Equal(t, tt.creator, creator)
+		})
+	}
+}
+
 func TestParseStackErrors(t *testing.T) {
 	tests := []struct {
 		name    string