From c832821064afc84c8a631a271892dd0c16fa2ae2 Mon Sep 17 00:00:00 2001
From: Eric Daniels <eric@erdaniels.com>
Date: Thu, 3 Oct 2024 09:48:49 -0400
Subject: [PATCH] Add tracebackancestors information to full stack

---
 internal/stack/stacks.go      | 30 ++++++++----------------------
 internal/stack/stacks_test.go | 27 ++++++++++++++++++++++++---
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/internal/stack/stacks.go b/internal/stack/stacks.go
index 1d31288..25a1b27 100644
--- a/internal/stack/stacks.go
+++ b/internal/stack/stacks.go
@@ -135,8 +135,9 @@ func (p *stackParser) parseStack(line string) (Stack, error) {
 
 	// Read the rest of the stack trace.
 	var (
-		firstFunction string
-		fullStack     bytes.Buffer
+		firstFunction    string
+		fullStack        bytes.Buffer
+		parsingAncestors bool
 	)
 	funcs := make(map[string]struct{})
 	for p.scan.Scan() {
@@ -152,7 +153,11 @@ func (p *stackParser) parseStack(line string) (Stack, error) {
 		fullStack.WriteString(line)
 		fullStack.WriteByte('\n') // scanner trims the newline
 
-		if len(line) == 0 {
+		if strings.HasPrefix(line, "[originating from goroutine ") {
+			parsingAncestors = true
+		}
+
+		if len(line) == 0 || parsingAncestors {
 			// Empty line usually marks the end of the stack
 			// but we don't want to have to rely on that.
 			// Just skip it.
@@ -201,25 +206,6 @@ func (p *stackParser) parseStack(line string) (Stack, error) {
 				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{
diff --git a/internal/stack/stacks_test.go b/internal/stack/stacks_test.go
index 4a8aebe..2dbd462 100644
--- a/internal/stack/stacks_test.go
+++ b/internal/stack/stacks_test.go
@@ -26,6 +26,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
 	"testing"
@@ -108,9 +109,20 @@ func TestCurrent(t *testing.T) {
 	assert.Contains(t, all, "stack/stacks_test.go",
 		"file name missing in stack:\n%s", all)
 
+	stackDepthFactor := 1
+	for _, debugVar := range strings.Split(os.Getenv("GODEBUG"), ",") {
+		debugKeyVal := strings.Split(debugVar, "=")
+		if len(debugKeyVal) == 2 && debugKeyVal[0] == "tracebackancestors" {
+			depth, err := strconv.ParseInt(debugKeyVal[1], 10, 32)
+			if err == nil && depth > 0 {
+				stackDepthFactor = int(depth) + 1
+			}
+		}
+	}
+
 	// Ensure that we are not returning the buffer without slicing it
 	// from getStackBuffer.
-	if len(got.Full()) > 1024 {
+	if len(got.Full()) > 1024*stackDepthFactor {
 		t.Fatalf("Returned stack is too large")
 	}
 }
@@ -355,8 +367,9 @@ func TestParseStackFixtures(t *testing.T) {
 		State         string
 		FirstFunction string
 
-		HasFunctions    []string // non-exhaustive, in any order
-		NotHasFunctions []string
+		HasFunctions      []string // non-exhaustive, in any order
+		NotHasFunctions   []string
+		HasFullStackLines []string // non-exhaustive, in any order
 	}
 
 	tests := []struct {
@@ -489,6 +502,9 @@ func TestParseStackFixtures(t *testing.T) {
 						"main.start",
 						"main.main",
 					},
+					HasFullStackLines: []string{
+						"	/usr/lib/go/src/net/http/transport.go:1437 +0x3cb",
+					},
 				},
 				{
 					ID:            4,
@@ -557,6 +573,11 @@ func TestParseStackFixtures(t *testing.T) {
 				for _, fn := range wantStack.NotHasFunctions {
 					assert.False(t, gotStack.HasFunction(fn), "unexpected in stack: %v\n%s", fn, gotStack.Full())
 				}
+
+				fullStack := gotStack.Full()
+				for _, line := range wantStack.HasFullStackLines {
+					assert.Contains(t, fullStack, line, "missing in full stack: %v\n%s", line, fullStack)
+				}
 			}
 
 			for _, s := range stacksByID {