Fix truncated Go CPU profiles (#3344)

* Fix truncated CPU profiles

* Write a test for minGroupSize.

Not quite sure where the requirement for minGroupSize 2 comes from

* Update maxRecursionDepth to 56

* Rename method, to something better maybe?

Author: simonswine

pkg/pprof/fix_go_heap_truncated_test.go

@@ -19,8 +19,8 @@ func FixGoProfile(p *profilev1.Profile) *profilev1.Profile {
 	// Now that the profile is normalized, we can try to repair
 	// truncated stack traces, if any. Note that repaired stacks
 	// are not deduplicated, so the caller need to normalize the
-	if MayHaveGoHeapTruncatedStacktraces(p) {
-		RepairGoHeapTruncatedStacktraces(p)
+	if PotentialTruncatedGoStacktraces(p) {
+		RepairGoTruncatedStacktraces(p)
 	}
 	return p
 }

pkg/pprof/fix_go_profile.go

@@ -10,7 +10,7 @@ import (
 )
 
 func Test_FixGoProfile(t *testing.T) {
-	p, err := OpenFile("testdata/goheapfix/heap_go_truncated_4.pb.gz")
+	p, err := OpenFile("testdata/gotruncatefix/heap_go_truncated_4.pb.gz")
 	require.NoError(t, err)
 
 	f := FixGoProfile(p.Profile)

pkg/pprof/fix_go_profile_test.go

@@ -10,30 +10,49 @@ import (
 )
 
 const (
-	minGroupSize = 2
+	minGroupSize      = 2
+	maxRecursiveDepth = 56 // Profiles with deeply recursive stack traces are ignored.
 
 	tokens    = 8
 	tokenLen  = 16
-	suffixLen = tokens + tokenLen
+	suffixLen = tokens + tokenLen // stacktraces shorter than suffixLen are not considered as truncated or missing stacks copied from
 
 	tokenBytesLen  = tokenLen * 8
 	suffixBytesLen = suffixLen * 8
 )
 
-// MayHaveGoHeapTruncatedStacktraces reports whether there are
+// PotentialTruncatedGoStacktraces reports whether there are
 // any chances that the profile may have truncated stack traces.
-func MayHaveGoHeapTruncatedStacktraces(p *profilev1.Profile) bool {
-	if !hasGoHeapSampleTypes(p) {
+func PotentialTruncatedGoStacktraces(p *profilev1.Profile) bool {
+	var minDepth int
+	var maxDepth int
+
+	if hasGoHeapSampleTypes(p) {
+		minDepth = 32
+
+		// Go heap profiles in Go 1.23+ have a depth limit of 128 frames. Let's not try to fix truncation if we see any longer stacks.
+		maxDepth = 33
+	} else if hasGoCPUSampleTypes(p) {
+		minDepth = 64
+	} else {
 		return false
 	}
-	// Some truncated stacks have depth less than the depth limit (32).
-	const minDepth = 28
+
+	// Some truncated heap stacks have depth less than the depth limit.
+	// https://github.com/golang/go/blob/f7c330eac7777612574d8a1652fd415391f6095e/src/runtime/mprof.go#L446
+	minDepth -= 4
+
+	deepEnough := false
 	for _, s := range p.Sample {
 		if len(s.LocationId) >= minDepth {
-			return true
+			deepEnough = true
+		}
+		// when it's too deep we no longer perform reassemlbing of truncated stacktraces
+		if maxDepth != 0 && len(s.LocationId) >= maxDepth {
+			return false
 		}
 	}
-	return false
+	return deepEnough
 }
 
 func hasGoHeapSampleTypes(p *profilev1.Profile) bool {
@@ -50,7 +69,18 @@ func hasGoHeapSampleTypes(p *profilev1.Profile) bool {
 	return false
 }
 
-// RepairGoHeapTruncatedStacktraces repairs truncated stack traces
+func hasGoCPUSampleTypes(p *profilev1.Profile) bool {
+	for _, st := range p.SampleType {
+		switch p.StringTable[st.Type] {
+		case
+			"cpu":
+			return true
+		}
+	}
+	return false
+}
+
+// RepairGoTruncatedStacktraces repairs truncated stack traces
 // in Go heap profiles.
 //
 // Go heap profile has a depth limit of 32 frames, which often
@@ -66,7 +96,7 @@ func hasGoHeapSampleTypes(p *profilev1.Profile) bool {
 // one by at least 16 frames in a row from the top, and has frames
 // above its root, stack S considered truncated, and the missing part
 // is copied from R.
-func RepairGoHeapTruncatedStacktraces(p *profilev1.Profile) {
+func RepairGoTruncatedStacktraces(p *profilev1.Profile) {
 	// Group stack traces by bottom (closest to the root) locations.
 	// Typically, there are very few groups (a hundred or two).
 	samples, groups := split(p)
@@ -82,7 +112,7 @@ func RepairGoHeapTruncatedStacktraces(p *profilev1.Profile) {
 		if i+1 < len(groups) {
 			n = groups[i+1]
 		}
-		if s := n - g; s < minGroupSize {
+		if s := n - g; s < (minGroupSize - 1) {
 			continue
 		}
 		// We take suffix of the first sample in the group.
@@ -168,8 +198,7 @@ func RepairGoHeapTruncatedStacktraces(p *profilev1.Profile) {
 			}
 			c = n
 			j++
-			if j == tokenLen {
-				// Profiles with deeply recursive stack traces are ignored.
+			if j == maxRecursiveDepth {
 				return
 			}
 		}

pkg/pprof/fix_go_heap_truncated.go renamed to pkg/pprof/fix_go_truncated.go

@@ -0,0 +1,89 @@
+package pprof
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/pyroscope/pkg/pprof/testhelper"
+)
+
+func Benchmark_RepairGoTruncatedStacktraces(b *testing.B) {
+	p, err := OpenFile("testdata/gotruncatefix/heap_go_truncated_3.pb.gz")
+	require.NoError(b, err)
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		RepairGoTruncatedStacktraces(FixGoProfile(p.CloneVT()))
+	}
+}
+
+func generateStackTrace(n int) []string {
+	res := make([]string, n)
+	runes := []rune("abcdefghijklmnopqrstuvwxyz")
+
+	for idx := range res {
+		dest := n - (idx + 1)
+		if idx == 0 {
+			res[dest] = "start"
+			continue
+		}
+		res[dest] = fmt.Sprintf("%c%d", runes[(idx-1)/10], (idx-1)%10)
+	}
+	return res
+}
+
+func Test_RepairGoTruncatedStacktraces(t *testing.T) {
+	n := 128
+	fullStack := generateStackTrace(n)
+	b := testhelper.NewProfileBuilder(0).CPUProfile()
+	b.ForStacktraceString(fullStack[n-24:]...).AddSamples(1)
+	b.ForStacktraceString(fullStack[n-58 : n-9]...).AddSamples(2)
+	b.ForStacktraceString(fullStack[n-57 : n-8]...).AddSamples(4)
+	b.ForStacktraceString(fullStack[n-56 : n-7]...).AddSamples(8)
+	b.ForStacktraceString(append([]string{"yy1"}, fullStack[n-22:]...)...).AddSamples(16)
+
+	RepairGoTruncatedStacktraces(b.Profile)
+
+	// ensure all stacktraces start with the same 8 location ids
+	stacks := make([]uint64, 8)
+	for idx, sample := range b.Profile.Sample {
+		first8Stacks := sample.LocationId[len(sample.LocationId)-8:]
+		if idx == 0 {
+			copy(stacks, first8Stacks)
+			continue
+		}
+		t.Log(stacks)
+		assert.Equal(t, stacks, first8Stacks)
+	}
+}
+
+func Test_UpdateFixtures_RepairGoTruncatedStacktraces(t *testing.T) {
+	if os.Getenv("UPDATE_FIXTURES") != "true" {
+		t.Skip()
+	}
+	t.Helper()
+	paths := []string{
+		"testdata/gotruncatefix/heap_go_truncated_1.pb.gz", // Cortex.
+		"testdata/gotruncatefix/heap_go_truncated_2.pb.gz", // Cortex.
+		"testdata/gotruncatefix/heap_go_truncated_3.pb.gz", // Loki. Pathological.
+		"testdata/gotruncatefix/heap_go_truncated_4.pb.gz", // Pyroscope.
+		"testdata/gotruncatefix/cpu_go_truncated_1.pb.gz",  // Cloudwatch Exporter
+	}
+	for _, path := range paths {
+		func() {
+			p, err := OpenFile(path)
+			require.NoError(t, err, path)
+			f, err := os.Create(path + ".fixed")
+			require.NoError(t, err, path)
+			defer f.Close()
+			p.Profile = FixGoProfile(p.Profile)
+			RepairGoTruncatedStacktraces(p.Profile)
+			_, err = p.WriteTo(f)
+			require.NoError(t, err, path)
+		}()
+	}
+}