fix: stack trace selector ignores conflicting functions (#3414)

Author: kolesnikovae

pkg/phlaredb/symdb/resolver_pprof.go

@@ -39,7 +39,7 @@ func buildPprof(
 	// 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, callSite: selection.callSite}
+		b = &pprofTree{maxNodes: maxNodes, selection: selection}
 	}
 	b.init(symbols, samples)
 	if err := symbols.Stacktraces.ResolveStacktraceLocations(ctx, b, samples.StacktraceIDs); err != nil {

pkg/phlaredb/symdb/resolver_pprof_test.go

@@ -141,6 +141,79 @@ func Test_Pprof_subtree(t *testing.T) {
 	require.Equal(t, expected, actual)
 }
 
+func Test_Pprof_subtree_multiple_versions(t *testing.T) {
+	profile := &googlev1.Profile{
+		StringTable: []string{"", "a", "b", "c", "d"},
+		Function: []*googlev1.Function{
+			{Id: 1, Name: 1},               // a
+			{Id: 2, Name: 2},               // b
+			{Id: 3, Name: 3},               // c
+			{Id: 4, Name: 4, StartLine: 1}, // d
+			{Id: 5, Name: 4, StartLine: 2}, // d(2)
+		},
+		Mapping: []*googlev1.Mapping{{Id: 1}},
+		Location: []*googlev1.Location{
+			{Id: 1, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 1, Line: 1}}}, // a
+			{Id: 2, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 1}}}, // b:1
+			{Id: 3, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 2}}}, // b:2
+			{Id: 4, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 3, Line: 1}}}, // c
+			{Id: 5, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 4, Line: 1}}}, // d
+			{Id: 6, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 5, Line: 1}}}, // d(2)
+		},
+		Sample: []*googlev1.Sample{
+			{LocationId: []uint64{5, 4, 2, 1}, Value: []int64{1}}, // a, b:1, c, d
+			{LocationId: []uint64{6, 4, 3, 1}, Value: []int64{1}}, // a, b:2, c, d(2)
+			{LocationId: []uint64{3, 1}, Value: []int64{1}},       // a, b:2
+			{LocationId: []uint64{4, 1}, Value: []int64{1}},       // a, c
+			{LocationId: []uint64{5}, Value: []int64{1}},          // d
+			{LocationId: []uint64{6}, Value: []int64{1}},          // d (2)
+		},
+	}
+
+	db := NewSymDB(DefaultConfig().WithDirectory(t.TempDir()))
+	w := db.WriteProfileSymbols(0, profile)
+	r := NewResolver(context.Background(), db,
+		WithResolverStackTraceSelector(&typesv1.StackTraceSelector{
+			CallSite: []*typesv1.Location{{Name: "a"}, {Name: "b"}, {Name: "c"}, {Name: "d"}},
+		}))
+
+	r.AddSamples(0, w[0].Samples)
+	actual, err := r.Pprof()
+	require.NoError(t, err)
+	// Sample order is not deterministic.
+	sort.Slice(actual.Sample, func(i, j int) bool {
+		return slices.Compare(actual.Sample[i].LocationId, actual.Sample[j].LocationId) >= 0
+	})
+
+	expected := &googlev1.Profile{
+		PeriodType:  &googlev1.ValueType{},
+		SampleType:  []*googlev1.ValueType{{}},
+		StringTable: []string{"", "a", "b", "c", "d"},
+		Function: []*googlev1.Function{
+			{Id: 1, Name: 1},               // a
+			{Id: 2, Name: 2},               // b
+			{Id: 3, Name: 3},               // c
+			{Id: 4, Name: 4, StartLine: 1}, // d
+			{Id: 5, Name: 4, StartLine: 2}, // d(2)
+		},
+		Mapping: []*googlev1.Mapping{{Id: 1}},
+		Location: []*googlev1.Location{
+			{Id: 1, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 1, Line: 1}}}, // a
+			{Id: 2, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 1}}}, // b:1
+			{Id: 3, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 2}}}, // b:2
+			{Id: 4, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 3, Line: 1}}}, // c
+			{Id: 5, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 4, Line: 1}}}, // d
+			{Id: 6, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 5, Line: 1}}}, // d(2)
+		},
+		Sample: []*googlev1.Sample{
+			{LocationId: []uint64{6, 4, 3, 1}, Value: []int64{1}}, // a, b:2, c, d(2)
+			{LocationId: []uint64{5, 4, 2, 1}, Value: []int64{1}}, // a, b:1, c, d
+		},
+	}
+
+	require.Equal(t, expected, actual)
+}
+
 func Test_Resolver_pprof_options(t *testing.T) {
 	s := newMemSuite(t, [][]string{{"testdata/profile.pb.gz"}})
 	samples := s.indexed[0][0].Samples

pkg/phlaredb/symdb/resolver_pprof_tree.go

@@ -31,8 +31,8 @@ type pprofTree struct {
 	functionsBuf []int32
 	locationsBuf []uint64
 
-	callSite []uint32
-	fnNames  func(locations []int32) ([]int32, bool)
+	selection *SelectedStackTraces
+	fnNames   func(locations []int32) ([]int32, bool)
 
 	// After truncation many samples will have the same stack trace.
 	// The map is used to deduplicate them. The key is sample.LocationId
@@ -57,7 +57,7 @@ func (r *pprofTree) init(symbols *Symbols, samples schemav1.Samples) {
 	r.functionTree = model.NewStacktraceTree(samples.Len() * 2)
 	r.stacktraces = make([]truncatedStacktraceSample, 0, samples.Len())
 	r.sampleMap = make(map[string]*googlev1.Sample, samples.Len())
-	if len(r.callSite) > 0 {
+	if r.selection != nil && len(r.selection.callSite) > 0 {
 		r.fnNames = r.locFunctionsFiltered
 	} else {
 		r.fnNames = r.locFunctions
@@ -92,15 +92,15 @@ func (r *pprofTree) locFunctions(locations []int32) ([]int32, bool) {
 func (r *pprofTree) locFunctionsFiltered(locations []int32) ([]int32, bool) {
 	r.functionsBuf = r.functionsBuf[:0]
 	var pos int
-	pathLen := len(r.callSite)
+	pathLen := int(r.selection.depth)
 	// Even if len(locations) < pathLen, we still
 	// need to inspect locations line by line.
 	for i := len(locations) - 1; i >= 0; i-- {
 		lines := r.symbols.Locations[locations[i]].Line
 		for j := len(lines) - 1; j >= 0; j-- {
 			f := lines[j].FunctionId
 			if pos < pathLen {
-				if r.callSite[pos] != f {
+				if r.selection.callSite[pos] != r.selection.funcNames[f] {
 					return nil, false
 				}
 				pos++

pkg/phlaredb/symdb/stacktrace_selection.go

@@ -44,13 +44,18 @@ type SelectedStackTraces struct {
 	symbols *Symbols
 	// Go PGO filter.
 	gopgo *typesv1.GoPGO
-	// call_site filter
+	// Call site filter
 	relations        map[uint32]stackTraceLocationRelation
 	callSiteSelector []*typesv1.Location
-	callSite         []uint32 // stack trace of the call site
-	location         uint32   // stack trace leaf
+	callSite         []string // call site strings in the original order.
+	location         string   // stack trace leaf function.
 	depth            uint32
 	buf              []uint64
+	// Function ID => name. The lookup table is used to
+	// avoid unnecessary indirect accesses through the
+	// strings[functions[id].Name] path. Instead, the
+	// name can be resolved directly funcNames[id].
+	funcNames []string
 }
 
 func SelectStackTraces(symbols *Symbols, selector *typesv1.StackTraceSelector) *SelectedStackTraces {
@@ -59,10 +64,14 @@ func SelectStackTraces(symbols *Symbols, selector *typesv1.StackTraceSelector) *
 		callSiteSelector: selector.GetCallSite(),
 		gopgo:            selector.GetGoPgo(),
 	}
-	x.callSite = findCallSite(symbols, x.callSiteSelector)
+	x.callSite = callSiteFunctions(x.callSiteSelector)
 	if x.depth = uint32(len(x.callSite)); x.depth > 0 {
 		x.location = x.callSite[x.depth-1]
 	}
+	x.funcNames = make([]string, len(symbols.Functions))
+	for i, f := range symbols.Functions {
+		x.funcNames[i] = symbols.Strings[f.Name]
+	}
 	return x
 }
 
@@ -144,64 +153,34 @@ func (x *SelectedStackTraces) appendStackTrace(locations []uint64) stackTraceLoc
 		lines := x.symbols.Locations[locations[i]].Line
 		for j := len(lines) - 1; j >= 0; j-- {
 			f := lines[j].FunctionId
-			n += eq(x.location, f)
-			if pos < x.depth && pos == l {
-				pos += eq(x.callSite[pos], f)
+			if x.location == x.funcNames[f] {
+				n++
+			}
+			if pos < x.depth && pos == l && x.callSite[pos] == x.funcNames[f] {
+				pos++
 			}
 			l++
 		}
 	}
 	if n == 0 {
 		return 0
 	}
+	var isLeaf uint32
 	leaf := x.symbols.Locations[locations[0]].Line[0]
-	isLeaf := eq(x.location, leaf.FunctionId)
-	inSubtree := ge(pos, x.depth)
-	return stackTraceLocationRelation(inSubtree | isLeaf<<1 | (1-isLeaf)<<2)
-}
-
-func eq(a, b uint32) uint32 {
-	if a == b {
-		return 1
+	if x.location == x.funcNames[leaf.FunctionId] {
+		isLeaf = 1
 	}
-	return 0
-}
-
-func ge(a, b uint32) uint32 {
-	if a >= b {
-		return 1
+	var inSubtree uint32
+	if pos >= x.depth {
+		inSubtree = 1
 	}
-	return 0
+	return stackTraceLocationRelation(inSubtree | isLeaf<<1 | (1-isLeaf)<<2)
 }
 
-// findCallSite returns the stack trace of the call site
-// where each element in the stack trace is represented by
-// the function ID. Call site is the last element.
-// TODO(kolesnikovae): Location should also include the line number.
-func findCallSite(symbols *Symbols, locations []*typesv1.Location) []uint32 {
-	if len(locations) == 0 {
-		return nil
-	}
-	m := make(map[string]uint32, len(locations))
-	for _, loc := range locations {
-		m[loc.Name] = 0
-	}
-	c := len(m) // Only count unique names.
-	for f := 0; f < len(symbols.Functions) && c > 0; f++ {
-		s := symbols.Strings[symbols.Functions[f].Name]
-		if _, ok := m[s]; ok {
-			// We assume that no functions have the same name.
-			// Otherwise, the last one takes precedence.
-			m[s] = uint32(f) // f is FunctionId
-			c--
-		}
-	}
-	if c > 0 {
-		return nil
-	}
-	callSite := make([]uint32, len(locations))
+func callSiteFunctions(locations []*typesv1.Location) []string {
+	callSite := make([]string, len(locations))
 	for i, loc := range locations {
-		callSite[i] = m[loc.Name]
+		callSite[i] = loc.Name
 	}
 	return callSite
 }