perf: pprof truncation without symbols

pkg/phlaredb/symdb/resolver.go

@@ -41,14 +41,6 @@ type Resolver struct {
 
 type ResolverOption func(*Resolver)
 
-// WithResolverMaxConcurrent specifies how many partitions
-// can be resolved concurrently.
-func WithResolverMaxConcurrent(n int) ResolverOption {
-	return func(r *Resolver) {
-		r.c = n
-	}
-}
-
 // WithResolverMaxNodes specifies the desired maximum number
 // of nodes the resulting profile should include.
 func WithResolverMaxNodes(n int64) ResolverOption {

pkg/phlaredb/symdb/resolver_pprof.go

@@ -2,6 +2,7 @@ package symdb
 
 import (
 	"context"
+	"unsafe"
 
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
@@ -35,11 +36,18 @@ func buildPprof(
 	// profile can exist. Otherwise, build an empty profile.
 	case !selection.HasValidCallSite():
 		return b.buildPprof(), nil
-	// Truncation is applicable when there is an explicit
-	// limit on the number of the nodes in the profile, or
-	// if stack traces should be filtered by the call site.
-	case maxNodes > 0 || len(selection.callSite) > 0:
-		b = &pprofTree{maxNodes: maxNodes, selection: selection}
+	// Stack trace filtering is only possible when the profile
+	// has functions (symbolized); for that, we first build a
+	// function call tree and then trim nodes according to the
+	// max nodes limit.
+	case len(selection.callSite) > 0:
+		b = &pprofFuncTree{maxNodes: maxNodes, selection: selection}
+	// Otherwise, if the max nodes limit is provided, we rely on
+	// the location tree, and ignore symbols altogether. Note that
+	// the result of truncation may be slightly different compared
+	// to the function tree.
+	case maxNodes > 0:
+		b = &pprofLocTree{maxNodes: maxNodes}
 	}
 	b.init(symbols, samples)
 	if err := symbols.Stacktraces.ResolveStacktraceLocations(ctx, b, samples.StacktraceIDs); err != nil {
@@ -261,3 +269,17 @@ func copyStrings(profile *googlev1.Profile, symbols *Symbols, lut []uint32) {
 		f.SystemName = int64(lut[f.SystemName+o])
 	}
 }
+
+func uint64sliceString(u []uint64) string {
+	if len(u) == 0 {
+		return ""
+	}
+	return unsafe.String((*byte)(unsafe.Pointer(&u[0])), len(u)*8)
+}
+
+func int32sliceString(u []int32) string {
+	if len(u) == 0 {
+		return ""
+	}
+	return unsafe.String((*byte)(unsafe.Pointer(&u[0])), len(u)*4)
+}

pkg/phlaredb/symdb/resolver_pprof_go_pgo.go

@@ -2,7 +2,6 @@ package symdb
 
 import (
 	"strings"
-	"unsafe"
 
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
@@ -57,10 +56,6 @@ func (r *pprofGoPGO) InsertStacktrace(_ uint32, locations []int32) {
 	r.cur++
 }
 
-func int32sliceString(u []int32) string {
-	return unsafe.String((*byte)(unsafe.Pointer(&u[0])), len(u)*4)
-}
-
 func (r *pprofGoPGO) buildPprof() *googlev1.Profile {
 	createSampleTypeStub(&r.profile)
 	r.appendSamples()

pkg/phlaredb/symdb/resolver_pprof_test.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"os"
 	"slices"
 	"sort"
 	"testing"
@@ -12,6 +13,7 @@
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
 	v1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
+	"github.com/grafana/pyroscope/pkg/pprof"
 )
 
 func Test_memory_Resolver_ResolvePprof(t *testing.T) {
@@ -25,6 +27,19 @@
 	require.Equal(t, expectedFingerprint, pprofFingerprint(resolved, 0))
 }
 
+func Test_memory_Resolver_ResolvePprof_REMOVE_ME(t *testing.T) {
+	s := newMemSuite(t, [][]string{{"testdata/profile.pb.gz"}})
+
+	r := NewResolver(context.Background(), s.db, WithResolverMaxNodes(64))
+	defer r.Release()
+	r.AddSamples(0, s.indexed[0][0].Samples)
+	resolved, err := r.Pprof()
+	require.NoError(t, err)
+
+	b, err := pprof.Marshal(resolved, true)
+	assert.NoError(t, os.WriteFile("testdata/profile.pb.truncated", b, 0644))
+}
+
 func Test_block_Resolver_ResolvePprof_multiple_partitions(t *testing.T) {
 	s := newBlockSuite(t, [][]string{
 		{"testdata/profile.pb.gz"},

pkg/phlaredb/symdb/resolver_pprof_tree.go → pkg/phlaredb/symdb/resolver_pprof_tree_funcs.go

@@ -1,8 +1,6 @@
 package symdb
 
 import (
-	"unsafe"
-
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	"github.com/grafana/pyroscope/pkg/model"
 	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
@@ -14,7 +12,7 @@ const (
 	truncatedNodeName = "other"
 )
 
-type pprofTree struct {
+type pprofFuncTree struct {
 	symbols *Symbols
 	samples *schemav1.Samples
 	profile googlev1.Profile
@@ -44,12 +42,12 @@ type pprofTree struct {
 }
 
 type truncatedStacktraceSample struct {
-	stacktraceID    uint32
-	functionNodeIdx int32
-	value           int64
+	stacktraceID uint32
+	nodeIdx      int32
+	value        int64
 }
 
-func (r *pprofTree) init(symbols *Symbols, samples schemav1.Samples) {
+func (r *pprofFuncTree) init(symbols *Symbols, samples schemav1.Samples) {
 	r.symbols = symbols
 	r.samples = &samples
 	// We optimistically assume that each stacktrace has only
@@ -64,21 +62,21 @@ func (r *pprofTree) init(symbols *Symbols, samples schemav1.Samples) {
 	}
 }
 
-func (r *pprofTree) InsertStacktrace(stacktraceID uint32, locations []int32) {
+func (r *pprofFuncTree) InsertStacktrace(stacktraceID uint32, locations []int32) {
 	value := int64(r.samples.Values[r.cur])
 	r.cur++
 	functions, ok := r.fnNames(locations)
 	if ok {
 		functionNodeIdx := r.functionTree.Insert(functions, value)
 		r.stacktraces = append(r.stacktraces, truncatedStacktraceSample{
-			stacktraceID:    stacktraceID,
-			functionNodeIdx: functionNodeIdx,
-			value:           value,
+			stacktraceID: stacktraceID,
+			nodeIdx:      functionNodeIdx,
+			value:        value,
 		})
 	}
 }
 
-func (r *pprofTree) locFunctions(locations []int32) ([]int32, bool) {
+func (r *pprofFuncTree) locFunctions(locations []int32) ([]int32, bool) {
 	r.functionsBuf = r.functionsBuf[:0]
 	for i := 0; i < len(locations); i++ {
 		lines := r.symbols.Locations[locations[i]].Line
@@ -89,7 +87,7 @@ func (r *pprofTree) locFunctions(locations []int32) ([]int32, bool) {
 	return r.functionsBuf, true
 }
 
-func (r *pprofTree) locFunctionsFiltered(locations []int32) ([]int32, bool) {
+func (r *pprofFuncTree) locFunctionsFiltered(locations []int32) ([]int32, bool) {
 	r.functionsBuf = r.functionsBuf[:0]
 	var pos int
 	pathLen := int(r.selection.depth)
@@ -115,7 +113,7 @@ func (r *pprofTree) locFunctionsFiltered(locations []int32) ([]int32, bool) {
 	return r.functionsBuf, true
 }
 
-func (r *pprofTree) buildPprof() *googlev1.Profile {
+func (r *pprofFuncTree) buildPprof() *googlev1.Profile {
 	r.markNodesForTruncation()
 	for _, n := range r.stacktraces {
 		r.addSample(n)
@@ -132,7 +130,7 @@ func (r *pprofTree) buildPprof() *googlev1.Profile {
 	return &r.profile
 }
 
-func (r *pprofTree) markNodesForTruncation() {
+func (r *pprofFuncTree) markNodesForTruncation() {
 	minValue := r.functionTree.MinValue(r.maxNodes)
 	if minValue == 0 {
 		return
@@ -145,11 +143,11 @@ func (r *pprofTree) markNodesForTruncation() {
 	}
 }
 
-func (r *pprofTree) addSample(n truncatedStacktraceSample) {
+func (r *pprofFuncTree) addSample(n truncatedStacktraceSample) {
 	// Find the original stack trace and remove truncated
 	// locations based on the truncated functions.
 	var off int
-	r.functionsBuf, off = r.buildFunctionsStack(r.functionsBuf, n.functionNodeIdx)
+	r.functionsBuf, off = r.buildFunctionsStack(r.functionsBuf, n.nodeIdx)
 	if off < 0 {
 		// The stack has no functions without the truncation mark.
 		r.fullyTruncated += n.value
@@ -182,7 +180,7 @@ func (r *pprofTree) addSample(n truncatedStacktraceSample) {
 	r.sampleMap[uint64sliceString(locationsCopy)] = s
 }
 
-func (r *pprofTree) buildFunctionsStack(funcs []int32, idx int32) ([]int32, int) {
+func (r *pprofFuncTree) buildFunctionsStack(funcs []int32, idx int32) ([]int32, int) {
 	offset := -1
 	funcs = funcs[:0]
 	for i := idx; i > 0; i = r.functionTree.Nodes[i].Parent {
@@ -196,7 +194,7 @@ func (r *pprofTree) buildFunctionsStack(funcs []int32, idx int32) ([]int32, int)
 	return funcs, offset
 }
 
-func (r *pprofTree) createSamples() {
+func (r *pprofFuncTree) createSamples() {
 	samples := len(r.sampleMap)
 	r.profile.Sample = make([]*googlev1.Sample, samples, samples+1)
 	var i int
@@ -216,10 +214,7 @@ func truncateLocations(locations []uint64, functions []int32, offset int, symbol
 	f := len(functions)
 	l := len(locations)
 	for ; l > 0 && f >= offset; l-- {
-		location := symbols.Locations[locations[l-1]]
-		for j := len(location.Line) - 1; j >= 0; j-- {
-			f--
-		}
+		f -= len(symbols.Locations[locations[l-1]].Line)
 	}
 	if l > 0 {
 		locations[0] = truncationMark
@@ -228,15 +223,7 @@ func truncateLocations(locations []uint64, functions []int32, offset int, symbol
 	return locations[l:]
 }
 
-func uint64sliceString(u []uint64) string {
-	if len(u) == 0 {
-		return ""
-	}
-	p := (*byte)(unsafe.Pointer(&u[0]))
-	return unsafe.String(p, len(u)*8)
-}
-
-func (r *pprofTree) createStubSample() {
+func (r *pprofFuncTree) createStubSample() {
 	r.profile.Sample = append(r.profile.Sample, &googlev1.Sample{
 		LocationId: []uint64{truncationMark},
 		Value:      []int64{r.fullyTruncated},
@@ -261,7 +248,6 @@ func createLocationStub(profile *googlev1.Profile) {
 		SystemName: stubNodeNameIdx,
 	}
 	profile.Function = append(profile.Function, stubFn)
-	// in the case there is no mapping, we need to create one
 	if len(profile.Mapping) == 0 {
 		profile.Mapping = append(profile.Mapping, &googlev1.Mapping{Id: 1})
 	}