github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/symdb/stacktrace_selection_test.go (about)

     1  package symdb
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
    11  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    12  	"github.com/grafana/pyroscope/pkg/slices"
    13  )
    14  
    15  func Test_StackTraceFilter(t *testing.T) {
    16  	profile := &googlev1.Profile{
    17  		StringTable: []string{"", "foo", "bar", "baz", "qux"},
    18  		Function: []*googlev1.Function{
    19  			{Id: 1, Name: 1},
    20  			{Id: 2, Name: 2},
    21  			{Id: 3, Name: 3},
    22  			{Id: 4, Name: 4},
    23  		},
    24  		Mapping: []*googlev1.Mapping{{Id: 1}},
    25  		Location: []*googlev1.Location{
    26  			{Id: 1, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 1, Line: 1}}}, // foo
    27  			{Id: 2, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 1}}}, // bar:1
    28  			{Id: 3, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 2, Line: 2}}}, // bar:2
    29  			{Id: 4, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 3, Line: 1}}}, // baz
    30  			{Id: 5, MappingId: 1, Line: []*googlev1.Line{{FunctionId: 4, Line: 1}}}, // qux
    31  		},
    32  		Sample: []*googlev1.Sample{
    33  			{LocationId: []uint64{4, 2, 1}, Value: []int64{1}}, // foo, bar:1, baz
    34  			{LocationId: []uint64{3, 1}, Value: []int64{1}},    // foo, bar:2
    35  			{LocationId: []uint64{4, 1}, Value: []int64{1}},    // foo, baz
    36  			{LocationId: []uint64{5}, Value: []int64{1}},       // qux
    37  
    38  			{LocationId: []uint64{2}, Value: []int64{1}},    // bar:1
    39  			{LocationId: []uint64{1, 2}, Value: []int64{1}}, // bar:1, foo
    40  			{LocationId: []uint64{3}, Value: []int64{1}},    // bar:2
    41  			{LocationId: []uint64{1, 3}, Value: []int64{1}}, // bar:2, foo
    42  		},
    43  	}
    44  
    45  	db := NewSymDB(DefaultConfig().WithDirectory(t.TempDir()))
    46  	w := db.WriteProfileSymbols(0, profile)
    47  
    48  	p, err := db.Partition(context.Background(), 0)
    49  	require.NoError(t, err)
    50  	symbols := p.Symbols()
    51  
    52  	type testCase struct {
    53  		selector *typesv1.StackTraceSelector
    54  		expected CallSiteValues
    55  	}
    56  
    57  	testCases := []testCase{
    58  		{
    59  			selector: &typesv1.StackTraceSelector{
    60  				CallSite: []*typesv1.Location{{Name: "foo"}},
    61  			},
    62  			expected: CallSiteValues{
    63  				Flat:          0,
    64  				Total:         3,
    65  				LocationFlat:  2,
    66  				LocationTotal: 5,
    67  			},
    68  		},
    69  		{
    70  			selector: &typesv1.StackTraceSelector{
    71  				CallSite: []*typesv1.Location{{Name: "bar"}},
    72  			},
    73  			expected: CallSiteValues{
    74  				Flat:          2,
    75  				Total:         4,
    76  				LocationFlat:  3,
    77  				LocationTotal: 6,
    78  			},
    79  		},
    80  		{
    81  			selector: &typesv1.StackTraceSelector{
    82  				CallSite: []*typesv1.Location{{Name: "foo"}, {Name: "bar"}},
    83  			},
    84  			expected: CallSiteValues{
    85  				Flat:          1,
    86  				Total:         2,
    87  				LocationFlat:  3,
    88  				LocationTotal: 6,
    89  			},
    90  		},
    91  		{
    92  			selector: &typesv1.StackTraceSelector{
    93  				CallSite: []*typesv1.Location{{Name: "foo"}, {Name: "bar"}, {Name: "baz"}},
    94  			},
    95  			expected: CallSiteValues{
    96  				Flat:          1,
    97  				Total:         1,
    98  				LocationFlat:  2,
    99  				LocationTotal: 2,
   100  			},
   101  		},
   102  		{
   103  			selector: &typesv1.StackTraceSelector{
   104  				CallSite: []*typesv1.Location{{Name: "foo"}, {Name: "bar"}, {Name: "baz"}, {Name: "qux"}},
   105  			},
   106  			expected: CallSiteValues{
   107  				Flat:          0,
   108  				Total:         0,
   109  				LocationFlat:  1,
   110  				LocationTotal: 1,
   111  			},
   112  		},
   113  		{selector: &typesv1.StackTraceSelector{}},
   114  		{},
   115  	}
   116  
   117  	for _, tc := range testCases {
   118  		selection := SelectStackTraces(symbols, tc.selector)
   119  		var values CallSiteValues
   120  		selection.CallSiteValues(&values, w[0].Samples)
   121  		assert.Equal(t, tc.expected, values, "selector: %+v", tc.selector)
   122  	}
   123  }
   124  
   125  func Benchmark_StackTraceFilter(b *testing.B) {
   126  	s := memSuite{t: b, files: [][]string{{"testdata/big-profile.pb.gz"}}}
   127  	s.config = DefaultConfig().WithDirectory(b.TempDir())
   128  	s.init()
   129  	samples := s.indexed[0][0].Samples
   130  
   131  	prt, err := s.db.Partition(context.Background(), 0)
   132  	require.NoError(b, err)
   133  	symbols := prt.Symbols()
   134  
   135  	p := s.profiles[0]
   136  	stack := p.Sample[len(p.Sample)/3].LocationId
   137  	selector := buildStackTraceSelector(p, stack[len(stack)/10:])
   138  	var values CallSiteValues
   139  
   140  	b.ReportAllocs()
   141  	b.ResetTimer()
   142  
   143  	selection := SelectStackTraces(symbols, selector)
   144  	for i := 0; i < b.N; i++ {
   145  		selection.CallSiteValues(&values, samples)
   146  	}
   147  }
   148  
   149  func buildStackTraceSelector(p *googlev1.Profile, locs []uint64) *typesv1.StackTraceSelector {
   150  	var selector typesv1.StackTraceSelector
   151  	for _, n := range locs {
   152  		for _, l := range p.Location[n-1].Line {
   153  			fn := p.Function[l.FunctionId-1]
   154  			selector.CallSite = append(selector.CallSite, &typesv1.Location{
   155  				Name: p.StringTable[fn.Name],
   156  			})
   157  		}
   158  	}
   159  	slices.Reverse(selector.CallSite)
   160  	return &selector
   161  }